| 8 | from .fMSE import MaskWeightedMSE |
| 9 | |
| 10 | class HDNetModel(BaseModel): |
| 11 | def __init__(self, opt): |
| 12 | BaseModel.__init__(self, opt) |
| 13 | # specify the training losses you want to print out. The training/test scripts will call <BaseModel.get_current_losses> |
| 14 | self.loss_names = ['G_L1'] |
| 15 | # specify the images you want to save/display. The training/test scripts will call <BaseModel.get_current_visuals> |
| 16 | self.visual_names = ['comp', 'real', 'output', 'mask', 'real_f', 'fake_f', 'bg', 'attentioned'] |
| 17 | # specify the models you want to save to the disk. The training/test scripts will call <BaseModel.save_networks> and <BaseModel.load_networks> |
| 18 | if self.isTrain: |
| 19 | self.model_names = ['G'] |
| 20 | else: |
| 21 | self.model_names = ['G'] |
| 22 | # define networks (both generator and discriminator) |
| 23 | self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf, opt.netG, opt.normG, |
| 24 | not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids) |
| 25 | self.relu = nn.ReLU() |
| 26 | if self.isTrain: |
| 27 | # define loss functions |
| 28 | self.criterionL1 = MaskWeightedMSE(100) |
| 29 | # initialize optimizers; schedulers will be automatically created by function <BaseModel.setup>. |
| 30 | self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr*opt.g_lr_ratio, betas=(opt.beta1, 0.999)) |
| 31 | self.optimizers.append(self.optimizer_G) |
| 32 | |
| 33 | def set_input(self, input): |
| 34 | """Unpack input data from the dataloader and perform necessary pre-processing steps. |
| 35 | Parameters: |
| 36 | input (dict): include the data itself and its metadata information. |
| 37 | """ |
| 38 | self.comp = input['comp'].to(self.device) |
| 39 | self.real = input['real'].to(self.device) |
| 40 | self.mask = input['mask'].to(self.device) |
| 41 | self.inputs = self.comp |
| 42 | if self.opt.input_nc == 4: |
| 43 | self.inputs = torch.cat([self.inputs, self.mask], 1) # channel-wise concatenation |
| 44 | self.real_f = self.real * self.mask |
| 45 | self.bg = self.real * (1 - self.mask) |
| 46 | |
| 47 | def forward(self): |
| 48 | self.output = self.netG(self.inputs, self.mask) |
| 49 | self.fake_f = self.output * self.mask |
| 50 | self.attentioned = self.output * self.mask + self.inputs[:,:3,:,:] * (1 - self.mask) |
| 51 | self.harmonized = self.attentioned |
| 52 | |
| 53 | def backward_G(self): |
| 54 | """Calculate GAN and L1 loss for the generator""" |
| 55 | self.loss_G_L1 = self.criterionL1(self.attentioned, self.real, self.mask) * self.opt.lambda_L1 |
| 56 | self.loss_G = self.loss_G_L1 |
| 57 | self.loss_G.backward() |
| 58 | |
| 59 | def optimize_parameters(self): |
| 60 | self.forward() |
| 61 | # update G |
| 62 | self.optimizer_G.zero_grad() # set G's gradients to zero |
| 63 | self.backward_G() # calculate graidents for G |
| 64 | self.optimizer_G.step() # udpate G's weights |
| 65 |
nothing calls this directly
no outgoing calls
no test coverage detected