RabbitRUI/ruispace
0
1"""This script defines the base network model for Deep3DFaceRecon_pytorch2"""3 4import os5import numpy as np6import torch7from collections import OrderedDict8from abc import ABC, abstractmethod9from . import networks10 11 12class BaseModel(ABC):13 """This class is an abstract base class (ABC) for models.14 To create a subclass, you need to implement the following five functions:15 -- <__init__>: initialize the class; first call BaseModel.__init__(self, opt).16 -- <set_input>: unpack data from dataset and apply preprocessing.17 -- <forward>: produce intermediate results.18 -- <optimize_parameters>: calculate losses, gradients, and update network weights.19 -- <modify_commandline_options>: (optionally) add model-specific options and set default options.20 """21 22 def __init__(self, opt):23 """Initialize the BaseModel class.24 25 Parameters:26 opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions27 28 When creating your custom class, you need to implement your own initialization.29 In this fucntion, you should first call <BaseModel.__init__(self, opt)>30 Then, you need to define four lists:31 -- self.loss_names (str list): specify the training losses that you want to plot and save.32 -- self.model_names (str list): specify the images that you want to display and save.33 -- self.visual_names (str list): define networks used in our training.34 -- self.optimizers (optimizer list): define and initialize optimizers. You can define one optimizer for each network. If two networks are updated at the same time, you can use itertools.chain to group them. See cycle_gan_model.py for an example.35 """36 self.opt = opt37 self.isTrain = False38 self.device = torch.device('cpu') 39 self.save_dir = " " # os.path.join(opt.checkpoints_dir, opt.name) # save all the checkpoints to save_dir40 self.loss_names = []41 self.model_names = []42 self.visual_names = []43 self.parallel_names = []44 self.optimizers = []45 self.image_paths = []46 self.metric = 0 # used for learning rate policy 'plateau'47 48 @staticmethod49 def dict_grad_hook_factory(add_func=lambda x: x):50 saved_dict = dict()51 52 def hook_gen(name):53 def grad_hook(grad):54 saved_vals = add_func(grad)55 saved_dict[name] = saved_vals56 return grad_hook57 return hook_gen, saved_dict58 59 @staticmethod60 def modify_commandline_options(parser, is_train):61 """Add new model-specific options, and rewrite default values for existing options.62 63 Parameters:64 parser -- original option parser65 is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.66 67 Returns:68 the modified parser.69 """70 return parser71 72 @abstractmethod73 def set_input(self, input):74 """Unpack input data from the dataloader and perform necessary pre-processing steps.75 76 Parameters:77 input (dict): includes the data itself and its metadata information.78 """79 pass80 81 @abstractmethod82 def forward(self):83 """Run forward pass; called by both functions <optimize_parameters> and <test>."""84 pass85 86 @abstractmethod87 def optimize_parameters(self):88 """Calculate losses, gradients, and update network weights; called in every training iteration"""89 pass90 91 def setup(self, opt):92 """Load and print networks; create schedulers93 94 Parameters:95 opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions96 """97 if self.isTrain:98 self.schedulers = [networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers]99 100 if not self.isTrain or opt.continue_train:101 load_suffix = opt.epoch102 self.load_networks(load_suffix)103 104 105 # self.print_networks(opt.verbose)106 107 def parallelize(self, convert_sync_batchnorm=True):108 if not self.opt.use_ddp:109 for name in self.parallel_names:110 if isinstance(name, str):111 module = getattr(self, name)112 setattr(self, name, module.to(self.device))113 else:114 for name in self.model_names:115 if isinstance(name, str):116 module = getattr(self, name)117 if convert_sync_batchnorm:118 module = torch.nn.SyncBatchNorm.convert_sync_batchnorm(module)119 setattr(self, name, torch.nn.parallel.DistributedDataParallel(module.to(self.device),120 device_ids=[self.device.index], 121 find_unused_parameters=True, broadcast_buffers=True))122 123 # DistributedDataParallel is not needed when a module doesn't have any parameter that requires a gradient.124 for name in self.parallel_names:125 if isinstance(name, str) and name not in self.model_names:126 module = getattr(self, name)127 setattr(self, name, module.to(self.device))128 129 # put state_dict of optimizer to gpu device130 if self.opt.phase != 'test':131 if self.opt.continue_train:132 for optim in self.optimizers:133 for state in optim.state.values():134 for k, v in state.items():135 if isinstance(v, torch.Tensor):136 state[k] = v.to(self.device)137 138 def data_dependent_initialize(self, data):139 pass140 141 def train(self):142 """Make models train mode"""143 for name in self.model_names:144 if isinstance(name, str):145 net = getattr(self, name)146 net.train()147 148 def eval(self):149 """Make models eval mode"""150 for name in self.model_names:151 if isinstance(name, str):152 net = getattr(self, name)153 net.eval()154 155 def test(self):156 """Forward function used in test time.157 158 This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop159 It also calls <compute_visuals> to produce additional visualization results160 """161 with torch.no_grad():162 self.forward()163 self.compute_visuals()164 165 def compute_visuals(self):166 """Calculate additional output images for visdom and HTML visualization"""167 pass168 169 def get_image_paths(self, name='A'):170 """ Return image paths that are used to load current data"""171 return self.image_paths if name =='A' else self.image_paths_B172 173 def update_learning_rate(self):174 """Update learning rates for all the networks; called at the end of every epoch"""175 for scheduler in self.schedulers:176 if self.opt.lr_policy == 'plateau':177 scheduler.step(self.metric)178 else:179 scheduler.step()180 181 lr = self.optimizers[0].param_groups[0]['lr']182 print('learning rate = %.7f' % lr)183 184 def get_current_visuals(self):185 """Return visualization images. train.py will display these images with visdom, and save the images to a HTML"""186 visual_ret = OrderedDict()187 for name in self.visual_names:188 if isinstance(name, str):189 visual_ret[name] = getattr(self, name)[:, :3, ...]190 return visual_ret191 192 def get_current_losses(self):193 """Return traning losses / errors. train.py will print out these errors on console, and save them to a file"""194 errors_ret = OrderedDict()195 for name in self.loss_names:196 if isinstance(name, str):197 errors_ret[name] = float(getattr(self, 'loss_' + name)) # float(...) works for both scalar tensor and float number198 return errors_ret199 200 def save_networks(self, epoch):201 """Save all the networks to the disk.202 203 Parameters:204 epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)205 """206 if not os.path.isdir(self.save_dir):207 os.makedirs(self.save_dir)208 209 save_filename = 'epoch_%s.pth' % (epoch)210 save_path = os.path.join(self.save_dir, save_filename)211 212 save_dict = {}213 for name in self.model_names:214 if isinstance(name, str):215 net = getattr(self, name)216 if isinstance(net, torch.nn.DataParallel) or isinstance(net,217 torch.nn.parallel.DistributedDataParallel):218 net = net.module219 save_dict[name] = net.state_dict()220 221 222 for i, optim in enumerate(self.optimizers):223 save_dict['opt_%02d'%i] = optim.state_dict()224 225 for i, sched in enumerate(self.schedulers):226 save_dict['sched_%02d'%i] = sched.state_dict()227 228 torch.save(save_dict, save_path)229 230 def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):231 """Fix InstanceNorm checkpoints incompatibility (prior to 0.4)"""232 key = keys[i]233 if i + 1 == len(keys): # at the end, pointing to a parameter/buffer234 if module.__class__.__name__.startswith('InstanceNorm') and \235 (key == 'running_mean' or key == 'running_var'):236 if getattr(module, key) is None:237 state_dict.pop('.'.join(keys))238 if module.__class__.__name__.startswith('InstanceNorm') and \239 (key == 'num_batches_tracked'):240 state_dict.pop('.'.join(keys))241 else:242 self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1)243 244 def load_networks(self, epoch):245 """Load all the networks from the disk.246 247 Parameters:248 epoch (int) -- current epoch; used in the file name '%s_net_%s.pth' % (epoch, name)249 """250 if self.opt.isTrain and self.opt.pretrained_name is not None:251 load_dir = os.path.join(self.opt.checkpoints_dir, self.opt.pretrained_name)252 else:253 load_dir = self.save_dir 254 load_filename = 'epoch_%s.pth' % (epoch)255 load_path = os.path.join(load_dir, load_filename)256 state_dict = torch.load(load_path, map_location=self.device)257 print('loading the model from %s' % load_path)258 259 for name in self.model_names:260 if isinstance(name, str):261 net = getattr(self, name)262 if isinstance(net, torch.nn.DataParallel):263 net = net.module264 net.load_state_dict(state_dict[name])265 266 if self.opt.phase != 'test':267 if self.opt.continue_train:268 print('loading the optim from %s' % load_path)269 for i, optim in enumerate(self.optimizers):270 optim.load_state_dict(state_dict['opt_%02d'%i])271 272 try:273 print('loading the sched from %s' % load_path)274 for i, sched in enumerate(self.schedulers):275 sched.load_state_dict(state_dict['sched_%02d'%i])276 except:277 print('Failed to load schedulers, set schedulers according to epoch count manually')278 for i, sched in enumerate(self.schedulers):279 sched.last_epoch = self.opt.epoch_count - 1280 281 282 283 284 def print_networks(self, verbose):285 """Print the total number of parameters in the network and (if verbose) network architecture286 287 Parameters:288 verbose (bool) -- if verbose: print the network architecture289 """290 print('---------- Networks initialized -------------')291 for name in self.model_names:292 if isinstance(name, str):293 net = getattr(self, name)294 num_params = 0295 for param in net.parameters():296 num_params += param.numel()297 if verbose:298 print(net)299 print('[Network %s] Total number of parameters : %.3f M' % (name, num_params / 1e6))300 print('-----------------------------------------------')301 302 def set_requires_grad(self, nets, requires_grad=False):303 """Set requies_grad=Fasle for all the networks to avoid unnecessary computations304 Parameters:305 nets (network list) -- a list of networks306 requires_grad (bool) -- whether the networks require gradients or not307 """308 if not isinstance(nets, list):309 nets = [nets]310 for net in nets:311 if net is not None:312 for param in net.parameters():313 param.requires_grad = requires_grad314 315 def generate_visuals_for_evaluation(self, data, mode):316 return {}317