mirror of https://github.com/kritiksoman/GIMP-ML
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
62 lines
1.8 KiB
Python
62 lines
1.8 KiB
Python
import os
|
|
import torch
|
|
|
|
|
|
class BaseModel():
|
|
def name(self):
|
|
return 'BaseModel'
|
|
|
|
def initialize(self, opt):
|
|
self.opt = opt
|
|
self.gpu_ids = opt.gpu_ids
|
|
self.isTrain = opt.isTrain
|
|
self.Tensor = torch.cuda.FloatTensor if (torch.cuda.is_available() and not opt.cFlag)else torch.Tensor
|
|
self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
|
|
|
|
def set_input(self, input):
|
|
self.input = input
|
|
|
|
def forward(self):
|
|
pass
|
|
|
|
# used in test time, no backprop
|
|
def test(self):
|
|
pass
|
|
|
|
def get_image_paths(self):
|
|
pass
|
|
|
|
def optimize_parameters(self):
|
|
pass
|
|
|
|
def get_current_visuals(self):
|
|
return self.input
|
|
|
|
def get_current_errors(self):
|
|
return {}
|
|
|
|
def save(self, label):
|
|
pass
|
|
|
|
# helper saving function that can be used by subclasses
|
|
def save_network(self, network, network_label, epoch_label, gpu_ids):
|
|
save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
|
|
save_path = os.path.join(self.save_dir, save_filename)
|
|
torch.save(network.cpu().state_dict(), save_path)
|
|
if len(gpu_ids) and torch.cuda.is_available():
|
|
network.cuda(device=gpu_ids[0])
|
|
|
|
# helper loading function that can be used by subclasses
|
|
def load_network(self, network, network_label, epoch_label):
|
|
save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
|
|
save_path = os.path.join(self.save_dir, save_filename)
|
|
if torch.cuda.is_available():
|
|
ckpt = torch.load(save_path)
|
|
else:
|
|
ckpt = torch.load(save_path, map_location=torch.device("cpu"))
|
|
ckpt = {key.replace("module.", ""): value for key, value in ckpt.items()}
|
|
network.load_state_dict(ckpt)
|
|
|
|
def update_learning_rate():
|
|
pass
|