SRGAN loss部分的pytorch代码实现

作者：雨丝儿

最近在参加华为与高校合做开发mindspore模型的活动，使用mindspore开发了SRGAN模型，下面几篇帖子想针对SRGAN做一些自己的经验分享。这篇帖子分享SRAGAN loss pytorch的实现。

Pytorch版本参考：https://github.com/dongheehand/SRGAN-PyTorch

Paper中SRGAN的loss:

对于Discriminator:

就是基础GAN中Discriminator的loss

代码实现:

其中gt为原始高分辨率图像，lr为gt经过双三次插值缩小四倍的低分辨率图像，cross_ent为BCELoss()

对与Generator:

Generator的loss包含三部分，一是基础的MSELoss,二是adversarial loss,三是将生成的HR图像与原始高清分辨率图像分别经过预训练的vgg19提取特征后，计算MSELoss.

代码部分：

VGG_loss = perceptual_loss(vgg_net)

cross_ent = nn.BCELoss()

tv_loss = TVLoss()

real_label = torch.ones((args.batch_size, 1)).to(device)

fake_label = torch.zeros((args.batch_size, 1)).to(device)

for i, tr_data in enumerate(loader):

gt = tr_data[‘GT’].to(device)

lr = tr_data[‘LR’].to(device)

output, _ = generator(lr)

fake_prob = discriminator(output)

# 第一部分

L2_loss = l2_loss(output, gt)

# 第二部分

adversarial_loss = args.adv_coeff * cross_ent(fake_prob, real_label)

# 第三部分

_percep_loss, hr_feat, sr_feat = VGG_loss((gt + 1.0) / 2.0, (output + 1.0) / 2.0, layer = args.feat_layer)

percep_loss = args.vgg_rescale_coeff * _percep_loss

g_loss = L2_loss + adversarial_loss + percep_loss

g_optim.zero_grad()

d_optim.zero_grad()

g_loss.backward()

g_optim.step()

其中vgg19是在imagenet上训练好的vgg19，选取其前37层，args.adv_coeff，args.vgg_rescale_coeff 为loss的系数，分别取0.003和0.006。

以上就是srgan loss部分的pytorch代码实现，下篇帖子将分享srgan loss部分minspore代码的实现。