s2f2e/base_model.py at master · klory/s2f2e · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
import torch
import os
import numpy as np

class BaseModel():
    def name(self, opt):
        return opt.model

    def initialize(self, opt):
        self.opt = opt
        self.isTrain = opt.isTrain
        self.save_dir = opt.save_dir
        self.Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.Tensor
        if not os.path.exists(self.save_dir):
            os.makedirs(self.save_dir)
        str_ids = self.opt.gpu_ids.split(',')
        self.opt.gpu_ids = []
        for str_id in str_ids:
            id = int(str_id)
            if id >= 0:
                self.opt.gpu_ids.append(id)

    def set_input(self):

        pass

    def forward(self):
        pass

    def test(self):
        pass

    def optimizer(self):
        pass

    def get_loss(self):
        pass

    def code_generate(self, label, batch_size):
        v1 = np.array([1] * 2 * 2).reshape(1, 1, 2, 2)
        v0 = np.array([0] * 2 * 2).reshape(1, 1, 2, 2)
        if label == 0:
            v = np.concatenate((v1, v0, v0), axis=1)
            v = np.repeat(v, batch_size, axis=0)
        elif label == 1:
            v = np.concatenate((v0, v1, v0), axis=1)
            v = np.repeat(v, batch_size, axis=0)
        if label == 2:
            v = np.concatenate((v0, v0, v1), axis=1)
            v = np.repeat(v, batch_size, axis=0)

        return self.Tensor(v.astype(np.float))

    def save_network(self, network, network_label, epoch_label):
        save_file = '%s_net_%s.pth' % (epoch_label, network_label)
        save_path = os.path.join(self.save_dir, save_file)
        torch.save(network.cpu().state_dict(), save_path)
        if torch.cuda.is_available():
            # TODO: ??? what is doing here
            network.cuda()

    def load_network(self, network, network_label, epoch_label):
        save_file = '%s_net_%s.pth' % (epoch_label, network_label)
        save_path = os.path.join(self.save_dir, save_file)
        network.load_state_dict(torch.load(save_path))