MLBench/ReaLens
0
1import os2import torch3import torch.distributed as dist4from pathlib import Path5from collections import OrderedDict6from abc import ABC, abstractmethod7from . import networks8 9 10class BaseModel(ABC):11 """This class is an abstract base class (ABC) for models.12 To create a subclass, you need to implement the following five functions:13 -- <__init__>: initialize the class; first call BaseModel.__init__(self, opt).14 -- <set_input>: unpack data from dataset and apply preprocessing.15 -- <forward>: produce intermediate results.16 -- <optimize_parameters>: calculate losses, gradients, and update network weights.17 -- <modify_commandline_options>: (optionally) add model-specific options and set default options.18 """19 20 def __init__(self, opt):21 """Initialize the BaseModel class.22 23 Parameters:24 opt (Option class)-- stores all the experiment flags; needs to be a subclass of BaseOptions25 26 When creating your custom class, you need to implement your own initialization.27 In this function, you should first call <BaseModel.__init__(self, opt)>28 Then, you need to define four lists:29 -- self.loss_names (str list): specify the training losses that you want to plot and save.30 -- self.model_names (str list): define networks used in our training.31 -- self.visual_names (str list): specify the images that you want to display and save.32 -- 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.33 """34 self.opt = opt35 self.isTrain = opt.isTrain36 self.save_dir = Path(opt.checkpoints_dir) / opt.name # save all the checkpoints to save_dir37 self.device = opt.device38 # with [scale_width], input images might have different sizes, which hurts the performance of cudnn.benchmark.39 if opt.preprocess != "scale_width":40 torch.backends.cudnn.benchmark = True41 self.loss_names = []42 self.model_names = []43 self.visual_names = []44 self.optimizers = []45 self.image_paths = []46 self.metric = 0 # used for learning rate policy 'plateau'47 48 @staticmethod49 def modify_commandline_options(parser, is_train):50 """Add new model-specific options, and rewrite default values for existing options.51 52 Parameters:53 parser -- original option parser54 is_train (bool) -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.55 56 Returns:57 the modified parser.58 """59 return parser60 61 @abstractmethod62 def set_input(self, input):63 """Unpack input data from the dataloader and perform necessary pre-processing steps.64 65 Parameters:66 input (dict): includes the data itself and its metadata information.67 """68 pass69 70 @abstractmethod71 def forward(self):72 """Run forward pass; called by both functions <optimize_parameters> and <test>."""73 pass74 75 @abstractmethod76 def optimize_parameters(self):77 """Calculate losses, gradients, and update network weights; called in every training iteration"""78 pass79 80 def setup(self, opt):81 """Load and print networks; create schedulers82 83 Parameters:84 opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions85 """86 # Initialize all networks and load if needed87 for name in self.model_names:88 if isinstance(name, str):89 net = getattr(self, "net" + name)90 net = networks.init_net(net, opt.init_type, opt.init_gain)91 92 # Load networks if needed93 if not self.isTrain or opt.continue_train:94 load_suffix = f"iter_{opt.load_iter}" if opt.load_iter > 0 else opt.epoch95 load_filename = f"{load_suffix}_net_{name}.pth"96 load_path = self.save_dir / load_filename97 98 if isinstance(net, torch.nn.parallel.DistributedDataParallel):99 net = net.module100 print(f"loading the model from {load_path}")101 102 state_dict = torch.load(load_path, map_location=str(self.device), weights_only=True)103 104 if hasattr(state_dict, "_metadata"):105 del state_dict._metadata106 107 # patch InstanceNorm checkpoints108 for key in list(state_dict.keys()):109 self.__patch_instance_norm_state_dict(state_dict, net, key.split("."))110 net.load_state_dict(state_dict)111 112 # Move network to device113 net.to(self.device)114 115 # Wrap networks with DDP after loading116 if dist.is_initialized():117 # Check if using syncbatch normalization for DDP118 if self.opt.norm == "syncbatch":119 raise ValueError(f"For distributed training, opt.norm must be 'syncbatch' or 'inst', but got '{self.opt.norm}'. " "Please set --norm syncbatch for multi-GPU training.")120 121 net = torch.nn.parallel.DistributedDataParallel(net, device_ids=[self.device.index])122 # Sync all processes after DDP wrapping123 dist.barrier()124 125 setattr(self, "net" + name, net)126 127 self.print_networks(opt.verbose)128 129 if self.isTrain:130 self.schedulers = [networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers]131 132 def eval(self):133 """Make models eval mode during test time"""134 for name in self.model_names:135 if isinstance(name, str):136 net = getattr(self, "net" + name)137 net.eval()138 139 def test(self):140 """Forward function used in test time.141 142 This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop143 It also calls <compute_visuals> to produce additional visualization results144 """145 with torch.no_grad():146 self.forward()147 self.compute_visuals()148 149 def compute_visuals(self):150 """Calculate additional output images for visdom and HTML visualization"""151 pass152 153 def get_image_paths(self):154 """Return image paths that are used to load current data"""155 return self.image_paths156 157 def update_learning_rate(self):158 """Update learning rates for all the networks; called at the end of every epoch"""159 old_lr = self.optimizers[0].param_groups[0]["lr"]160 for scheduler in self.schedulers:161 if self.opt.lr_policy == "plateau":162 scheduler.step(self.metric)163 else:164 scheduler.step()165 166 lr = self.optimizers[0].param_groups[0]["lr"]167 print(f"learning rate {old_lr:.7f} -> {lr:.7f}")168 169 def get_current_visuals(self):170 """Return visualization images. train.py will display these images with visdom, and save the images to a HTML"""171 visual_ret = OrderedDict()172 for name in self.visual_names:173 if isinstance(name, str):174 visual_ret[name] = getattr(self, name)175 return visual_ret176 177 def get_current_losses(self):178 """Return traning losses / errors. train.py will print out these errors on console, and save them to a file"""179 errors_ret = OrderedDict()180 for name in self.loss_names:181 if isinstance(name, str):182 errors_ret[name] = float(getattr(self, "loss_" + name)) # float(...) works for both scalar tensor and float number183 return errors_ret184 185 def save_networks(self, epoch):186 """Save all the networks to the disk, unwrapping them first."""187 188 # Only allow the main process (rank 0) to save the checkpoint189 if not dist.is_initialized() or dist.get_rank() == 0:190 for name in self.model_names:191 if isinstance(name, str):192 save_filename = f"{epoch}_net_{name}.pth"193 save_path = self.save_dir / save_filename194 net = getattr(self, "net" + name)195 196 # 1. First, unwrap from DDP if it exists197 if hasattr(net, "module"):198 model_to_save = net.module199 else:200 model_to_save = net201 202 # 2. Second, unwrap from torch.compile if it exists203 if hasattr(model_to_save, "_orig_mod"):204 model_to_save = model_to_save._orig_mod205 206 # 3. Save the final, clean state_dict207 torch.save(model_to_save.state_dict(), save_path)208 209 def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):210 """Fix InstanceNorm checkpoints incompatibility (prior to 0.4)"""211 key = keys[i]212 if i + 1 == len(keys): # at the end, pointing to a parameter/buffer213 if module.__class__.__name__.startswith("InstanceNorm") and (key == "running_mean" or key == "running_var"):214 if getattr(module, key) is None:215 state_dict.pop(".".join(keys))216 if module.__class__.__name__.startswith("InstanceNorm") and (key == "num_batches_tracked"):217 state_dict.pop(".".join(keys))218 else:219 self.__patch_instance_norm_state_dict(state_dict, getattr(module, key), keys, i + 1)220 221 def load_networks(self, epoch):222 """Load all networks from the disk for DDP."""223 224 for name in self.model_names:225 if isinstance(name, str):226 load_filename = f"{epoch}_net_{name}.pth"227 load_path = self.save_dir / load_filename228 net = getattr(self, "net" + name)229 230 if isinstance(net, torch.nn.parallel.DistributedDataParallel):231 net = net.module232 print(f"loading the model from {load_path}")233 234 state_dict = torch.load(load_path, map_location=str(self.device), weights_only=True)235 236 if hasattr(state_dict, "_metadata"):237 del state_dict._metadata238 239 # patch InstanceNorm checkpoints240 for key in list(state_dict.keys()):241 self.__patch_instance_norm_state_dict(state_dict, net, key.split("."))242 net.load_state_dict(state_dict)243 244 # Add a barrier to sync all processes before continuing245 if dist.is_initialized():246 dist.barrier()247 248 def print_networks(self, verbose):249 """Print the total number of parameters in the network and (if verbose) network architecture250 251 Parameters:252 verbose (bool) -- if verbose: print the network architecture253 """254 print("---------- Networks initialized -------------")255 for name in self.model_names:256 if isinstance(name, str):257 net = getattr(self, "net" + name)258 num_params = 0259 for param in net.parameters():260 num_params += param.numel()261 if verbose:262 print(net)263 print(f"[Network {name}] Total number of parameters : {num_params / 1e6:.3f} M")264 print("-----------------------------------------------")265 266 def set_requires_grad(self, nets, requires_grad=False):267 """Set requies_grad=Fasle for all the networks to avoid unnecessary computations268 Parameters:269 nets (network list) -- a list of networks270 requires_grad (bool) -- whether the networks require gradients or not271 """272 if not isinstance(nets, list):273 nets = [nets]274 for net in nets:275 if net is not None:276 for param in net.parameters():277 param.requires_grad = requires_grad278 279 def init_networks(self, init_type="normal", init_gain=0.02):280 """Initialize all networks: 1. move to device; 2. initialize weights281 282 Parameters:283 init_type (str) -- initialization method: normal | xavier | kaiming | orthogonal284 init_gain (float) -- scaling factor for normal, xavier and orthogonal285 """286 import os287 288 for name in self.model_names:289 if isinstance(name, str):290 net = getattr(self, "net" + name)291 292 # Move to device293 if torch.cuda.is_available():294 if "LOCAL_RANK" in os.environ:295 local_rank = int(os.environ["LOCAL_RANK"])296 net.to(local_rank)297 print(f"Initialized network {name} with device cuda:{local_rank}")298 else:299 net.to(0)300 print(f"Initialized network {name} with device cuda:0")301 else:302 net.to("cpu")303 print(f"Initialized network {name} with device cpu")304 305 # Initialize weights using networks function306 networks.init_weights(net, init_type, init_gain)307