Back propagation with customize loss function

username: Helen_R6
cooked: <p>Hi,</p>
<p>I am building a model for image synthesis. The architecture for the model is shown below.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/6/66e743ee9584b7a41fd79ec8af12de091cc1da73.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/66e743ee9584b7a41fd79ec8af12de091cc1da73" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/6/66e743ee9584b7a41fd79ec8af12de091cc1da73_2_690x265.jpeg" alt="image" data-base62-sha1="eGkbvXl7iNVpShwBG2OBaDlQ9Hl" width="690" height="265" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/6/66e743ee9584b7a41fd79ec8af12de091cc1da73_2_690x265.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/6/6/66e743ee9584b7a41fd79ec8af12de091cc1da73_2_1035x397.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/6/66e743ee9584b7a41fd79ec8af12de091cc1da73_2_1380x530.jpeg 2x" data-dominant-color="CDE6CD"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">2154×830 350 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
the objective is to update start_input to minimize the difference between start_output and target_output. (start_output and target_output are the activation for certain layer in ResNet )</p>
<p>the problem is that when I run the code below, the loss value doesn’t decrease, it is oscillating between certain value, and the distance between start_input and target_input is even increasing sometimes.<br>
Could any one help me figure out where went wrong? Any suggestions and help would be appreciated. Thanks in advance.</p>
<p>the code is here</p>
<pre><code class="lang-auto">
import torchvision.models as models
import torch 
import torch.nn as nn
import logging
from pytorch_pretrained_biggan import (BigGAN, one_hot_from_names, truncated_noise_sample,
                                       save_as_images, display_in_terminal)
import torch.nn.functional as F
import urllib
from PIL import Image
from torchvision import transforms
from numpy import linalg as LA

# load the pretrained weight for resnet and BigGAN
resnet = models.resnet50(pretrained=True)
GAN_model = BigGAN.from_pretrained('biggan-deep-256')

# define class for combine the two models
class MyEnsemble(nn.Module):
    def __init__(self, modelA, modelB):
        super(MyEnsemble, self).__init__()
        self.modelA=modelA
        self.modelB=modelB
    def forward(self,z,class_name, truncation):
        x=self.modelA(z, class_name, truncation)
        # resize the tensor 
        x1=F.interpolate(x,(224,224))
        x2=self.modelB(x1)
        return x1, x2 
truncation = 0.5
m = nn.Softmax(dim=1)
input = torch.randn(1, 1000)
class_vector = m(input)
modelA=GAN_model
modelB=resnet
model=MyEnsemble(modelA, modelB)


def my_loss(target_output,start_output):
    return torch.norm(target_output-start_output)

# generate random input for start point image
target_input=torch.randn(1,128,requires_grad = True)
target_point,target_output=model(target_input,class_vector,truncation)
target_output=target_output.detach()

# generate random input for start point image
start_input=torch.randn(1,128,requires_grad = True)

optimizer=torch.optim.Adam( [start_input,class_vector],lr=0.05)
loss_values=[]
distance_values=[]
model.eval()
for epoch in range(101):
    print (epoch)
    torch.autograd.set_detect_anomaly(True)
    optimizer.zero_grad()
    running_loss=0.0
    running_distance=0.0
    distance=torch.norm(start_input-target_input)
    start_point,start_output=model(start_input,class_vector,truncation)
    loss=my_loss(target_output,start_output)
    print(loss)
    print(distance)
    loss.backward(retain_graph=True)
    print('Epoch-{0} lr: {1}'.format(epoch, optimizer.param_groups[0]['lr']))
    optimizer.step()
    if epoch==100:
        save_as_images(start_point,file_name='result')
    running_loss+=loss.item()
    running_distance+=distance.item()
    loss_values.append(running_loss)
    distance_values.append(running_distance)
    


</code></pre>