Spaces:
Paused
Paused
| import torch | |
| from .base_model import BaseModel | |
| from . import model_utils | |
| class StarGANModel(BaseModel): | |
| """docstring for StarGANModel""" | |
| def __init__(self): | |
| super(StarGANModel, self).__init__() | |
| self.name = "StarGAN" | |
| def initialize(self, opt): | |
| super(StarGANModel, self).initialize(opt) | |
| self.net_gen = model_utils.define_splitG(self.opt.img_nc, self.opt.aus_nc, self.opt.ngf, use_dropout=self.opt.use_dropout, | |
| norm=self.opt.norm, init_type=self.opt.init_type, init_gain=self.opt.init_gain, gpu_ids=self.gpu_ids) | |
| self.models_name.append('gen') | |
| if self.is_train: | |
| self.net_dis = model_utils.define_splitD(self.opt.img_nc, self.opt.aus_nc, self.opt.final_size, self.opt.ndf, | |
| norm=self.opt.norm, init_type=self.opt.init_type, init_gain=self.opt.init_gain, gpu_ids=self.gpu_ids) | |
| self.models_name.append('dis') | |
| if self.opt.load_epoch > 0: | |
| self.load_ckpt(self.opt.load_epoch) | |
| def setup(self): | |
| super(StarGANModel, self).setup() | |
| if self.is_train: | |
| # setup optimizer | |
| self.optim_gen = torch.optim.Adam(self.net_gen.parameters(), | |
| lr=self.opt.lr, betas=(self.opt.beta1, 0.999)) | |
| self.optims.append(self.optim_gen) | |
| self.optim_dis = torch.optim.Adam(self.net_dis.parameters(), | |
| lr=self.opt.lr, betas=(self.opt.beta1, 0.999)) | |
| self.optims.append(self.optim_dis) | |
| # setup schedulers | |
| self.schedulers = [model_utils.get_scheduler(optim, self.opt) for optim in self.optims] | |
| def feed_batch(self, batch): | |
| self.src_img = batch['src_img'].to(self.device) | |
| self.tar_aus = batch['tar_aus'].type(torch.FloatTensor).to(self.device) | |
| if self.is_train: | |
| self.src_aus = batch['src_aus'].type(torch.FloatTensor).to(self.device) | |
| self.tar_img = batch['tar_img'].to(self.device) | |
| def forward(self): | |
| # generate fake image | |
| self.fake_img, _, _ = self.net_gen(self.src_img, self.tar_aus) | |
| # reconstruct real image | |
| if self.is_train: | |
| self.rec_real_img, _, _ = self.net_gen(self.fake_img, self.src_aus) | |
| def backward_dis(self): | |
| # real image | |
| pred_real, self.pred_real_aus = self.net_dis(self.src_img) | |
| self.loss_dis_real = self.criterionGAN(pred_real, True) | |
| self.loss_dis_real_aus = self.criterionMSE(self.pred_real_aus, self.src_aus) | |
| # fake image, detach to stop backward to generator | |
| pred_fake, _ = self.net_dis(self.fake_img.detach()) | |
| self.loss_dis_fake = self.criterionGAN(pred_fake, False) | |
| # combine dis loss | |
| self.loss_dis = self.opt.lambda_dis * (self.loss_dis_fake + self.loss_dis_real) \ | |
| + self.opt.lambda_aus * self.loss_dis_real_aus | |
| if self.opt.gan_type == 'wgan-gp': | |
| self.loss_dis_gp = self.gradient_penalty(self.src_img, self.fake_img) | |
| self.loss_dis = self.loss_dis + self.opt.lambda_wgan_gp * self.loss_dis_gp | |
| # backward discriminator loss | |
| self.loss_dis.backward() | |
| def backward_gen(self): | |
| # original to target domain, should fake the discriminator | |
| pred_fake, self.pred_fake_aus = self.net_dis(self.fake_img) | |
| self.loss_gen_GAN = self.criterionGAN(pred_fake, True) | |
| self.loss_gen_fake_aus = self.criterionMSE(self.pred_fake_aus, self.tar_aus) | |
| # target to original domain reconstruct, identity loss | |
| self.loss_gen_rec = self.criterionL1(self.rec_real_img, self.src_img) | |
| # combine and backward G loss | |
| self.loss_gen = self.opt.lambda_dis * self.loss_gen_GAN \ | |
| + self.opt.lambda_aus * self.loss_gen_fake_aus \ | |
| + self.opt.lambda_rec * self.loss_gen_rec | |
| self.loss_gen.backward() | |
| def optimize_paras(self, train_gen): | |
| self.forward() | |
| # update discriminator | |
| self.set_requires_grad(self.net_dis, True) | |
| self.optim_dis.zero_grad() | |
| self.backward_dis() | |
| self.optim_dis.step() | |
| # update G if needed | |
| if train_gen: | |
| self.set_requires_grad(self.net_dis, False) | |
| self.optim_gen.zero_grad() | |
| self.backward_gen() | |
| self.optim_gen.step() | |
| def save_ckpt(self, epoch): | |
| # save the specific networks | |
| save_models_name = ['gen', 'dis'] | |
| return super(StarGANModel, self).save_ckpt(epoch, save_models_name) | |
| def load_ckpt(self, epoch): | |
| # load the specific part of networks | |
| load_models_name = ['gen'] | |
| if self.is_train: | |
| load_models_name.extend(['dis']) | |
| return super(StarGANModel, self).load_ckpt(epoch, load_models_name) | |
| def clean_ckpt(self, epoch): | |
| # load the specific part of networks | |
| load_models_name = ['gen', 'dis'] | |
| return super(StarGANModel, self).clean_ckpt(epoch, load_models_name) | |
| def get_latest_losses(self): | |
| get_losses_name = ['dis_fake', 'dis_real', 'dis_real_aus', 'gen_rec'] | |
| return super(StarGANModel, self).get_latest_losses(get_losses_name) | |
| def get_latest_visuals(self): | |
| visuals_name = ['src_img', 'tar_img', 'fake_img'] | |
| if self.is_train: | |
| visuals_name.extend(['rec_real_img']) | |
| return super(StarGANModel, self).get_latest_visuals(visuals_name) | |