RabbitRUI/ruispace

1"""Model class template2 3This module provides a template for users to implement custom models.4You can specify '--model template' to use this model.5The class name should be consistent with both the filename and its model option.6The filename should be <model>_dataset.py7The class name should be <Model>Dataset.py8It implements a simple image-to-image translation baseline based on regression loss.9Given input-output pairs (data_A, data_B), it learns a network netG that can minimize the following L1 loss:10    min_<netG> ||netG(data_A) - data_B||_111You need to implement the following functions:12    <modify_commandline_options>:　Add model-specific options and rewrite default values for existing options.13    <__init__>: Initialize this model class.14    <set_input>: Unpack input data and perform data pre-processing.15    <forward>: Run forward pass. This will be called by both <optimize_parameters> and <test>.16    <optimize_parameters>: Update network weights; it will be called in every training iteration.17"""18import numpy as np19import torch20from .base_model import BaseModel21from . import networks22 23 24class TemplateModel(BaseModel):25    @staticmethod26    def modify_commandline_options(parser, is_train=True):27        """Add new model-specific options and rewrite default values for existing options.28 29        Parameters:30            parser -- the option parser31            is_train -- if it is training phase or test phase. You can use this flag to add training-specific or test-specific options.32 33        Returns:34            the modified parser.35        """36        parser.set_defaults(dataset_mode='aligned')  # You can rewrite default values for this model. For example, this model usually uses aligned dataset as its dataset.37        if is_train:38            parser.add_argument('--lambda_regression', type=float, default=1.0, help='weight for the regression loss')  # You can define new arguments for this model.39 40        return parser41 42    def __init__(self, opt):43        """Initialize this model class.44 45        Parameters:46            opt -- training/test options47 48        A few things can be done here.49        - (required) call the initialization function of BaseModel50        - define loss function, visualization images, model names, and optimizers51        """52        BaseModel.__init__(self, opt)  # call the initialization method of BaseModel53        # specify the training losses you want to print out. The program will call base_model.get_current_losses to plot the losses to the console and save them to the disk.54        self.loss_names = ['loss_G']55        # specify the images you want to save and display. The program will call base_model.get_current_visuals to save and display these images.56        self.visual_names = ['data_A', 'data_B', 'output']57        # specify the models you want to save to the disk. The program will call base_model.save_networks and base_model.load_networks to save and load networks.58        # you can use opt.isTrain to specify different behaviors for training and test. For example, some networks will not be used during test, and you don't need to load them.59        self.model_names = ['G']60        # define networks; you can use opt.isTrain to specify different behaviors for training and test.61        self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, gpu_ids=self.gpu_ids)62        if self.isTrain:  # only defined during training time63            # define your loss functions. You can use losses provided by torch.nn such as torch.nn.L1Loss.64            # We also provide a GANLoss class "networks.GANLoss". self.criterionGAN = networks.GANLoss().to(self.device)65            self.criterionLoss = torch.nn.L1Loss()66            # define and initialize optimizers. You can define one optimizer for each network.67            # 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.68            self.optimizer = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))69            self.optimizers = [self.optimizer]70 71        # Our program will automatically call <model.setup> to define schedulers, load networks, and print networks72 73    def set_input(self, input):74        """Unpack input data from the dataloader and perform necessary pre-processing steps.75 76        Parameters:77            input: a dictionary that contains the data itself and its metadata information.78        """79        AtoB = self.opt.direction == 'AtoB'  # use <direction> to swap data_A and data_B80        self.data_A = input['A' if AtoB else 'B'].to(self.device)  # get image data A81        self.data_B = input['B' if AtoB else 'A'].to(self.device)  # get image data B82        self.image_paths = input['A_paths' if AtoB else 'B_paths']  # get image paths83 84    def forward(self):85        """Run forward pass. This will be called by both functions <optimize_parameters> and <test>."""86        self.output = self.netG(self.data_A)  # generate output image given the input data_A87 88    def backward(self):89        """Calculate losses, gradients, and update network weights; called in every training iteration"""90        # caculate the intermediate results if necessary; here self.output has been computed during function <forward>91        # calculate loss given the input and intermediate results92        self.loss_G = self.criterionLoss(self.output, self.data_B) * self.opt.lambda_regression93        self.loss_G.backward()       # calculate gradients of network G w.r.t. loss_G94 95    def optimize_parameters(self):96        """Update network weights; it will be called in every training iteration."""97        self.forward()               # first call forward to calculate intermediate results98        self.optimizer.zero_grad()   # clear network G's existing gradients99        self.backward()              # calculate gradients for network G100        self.optimizer.step()        # update gradients for network G101