天天看點

感覺損失的pytorch實作

常用的損失函數如torch.nn.MSELoss()、torch.nn.L1Loss()計算的是兩幅圖像間像素位置一一對應的像素級的損失函數。對于一種極端情況:兩幅完全一樣的圖像但對應的像素位置錯開一個像素的距離,像素級損失函數計算出來的效果往往比較大,但人眼觀察兩幅圖像幾乎沒有差距。

感覺損失函數是一種能了解圖像語義資訊的損失函數,它利用一個Vgg網絡提取出圖像的特征資訊,并利用這些特征資訊做loss,能夠有效解決上述像素級損失函數的問題。以下是感覺損失的pytorch實作代碼:

定義感覺損失

# Loss functions
class PerceptualLoss():
	def contentFunc(self):
		conv_3_3_layer = 14
		cnn = models.vgg19(pretrained=True).features
		cnn = cnn.cuda()
		model = nn.Sequential()
		model = model.cuda()
		for i,layer in enumerate(list(cnn)):
			model.add_module(str(i),layer)
			if i == conv_3_3_layer:
				break
		return model
		
	def __init__(self, loss):
		self.criterion = loss
		self.contentFunc = self.contentFunc()
			
	def get_loss(self, fakeIm, realIm):
		f_fake = self.contentFunc.forward(fakeIm)
		f_real = self.contentFunc.forward(realIm)
		f_real_no_grad = f_real.detach()
		loss = self.criterion(f_fake, f_real_no_grad)
		return loss
           

感覺損失需要嵌套mse loss

感覺損失的pytorch實作
content_loss = PerceptualLoss(torch.nn.MSELoss())
loss_PerceptualLoss = content_loss.get_loss( fake_B, real_B)
loss_G = loss_GAN + lambda_A * loss_pixel + lambda_B * loss_lambda  #
           

繼續閱讀