Parameter not learning

Aniruth_Sundararaja1: <p>I made a 2 step  model with a u net and a gan as 2 consecutive steps.</p>
<p>the output that i get from the u net , i apply thresholding to get a mask , and pass the output and mask to the gan for inpainting.<br>
i want to make the threshold also learnable .<br>
i kept the threshold as nn.Parameter() , and also set required_grad = True , but then when I checkked while training the model , the parameter value is not getting updated at all.</p>
<p>The same init value of 0.5 is only coming.<br>
Code</p>
<p>class Combined_Model(nn.Module):</p>
<pre><code>def __init__(self , options):
    super(Combined_Model, self).__init__()

    self.pretrained_state_dict = torch.load(os.path.join(options.pretrained, 'G0000000.pt'), map_location=torch.device('cuda'))
    self.unet = UNet().to(options.device)


    if options.with_prompts:
        self.inpainter = Prompted_InpaintGenerator(options)
        self.org_gan = InpaintGenerator(options)
        #self.inpainter.load_state_dict(load_pretrained_weights(self.org_gan, self.pretrained_state_dict), strict=False)
        self.inpainter.load_state_dict(load_pretrained_weights(self.org_gan , self.inpainter) , strict=True)
    else:
        self.inpainter = InpaintGenerator(options)
        self.inpainter.load_state_dict(torch.load(os.path.join(options.pretrained, 'G0000000.pt'), map_location=options.device), strict=False)

    self.models = [self.unet, self.inpainter]

    self.learnable_threshold = nn.Parameter(torch.tensor(0.5), requires_grad=True)

def forward(self , x):

    unet_output = self.unet(x)
    unet_output_gray = tensor_to_cv2_gray(unet_output)
    flary_img_gray = tensor_to_cv2_gray(x)
    print(self.learnable_threshold)
    difference = (torch.from_numpy(flary_img_gray) - torch.from_numpy(unet_output_gray))
    #difference_tensor = torch.tensor(difference, dtype=torch.float32).to(options.device)
    difference_tensor = difference.clone().to(options.device)
    binary_mask = torch.where(difference_tensor &gt; self.learnable_threshold, torch.tensor(1.0).to(options.device), torch.tensor(0.0).to(options.device))
    binary_mask = binary_mask.unsqueeze(1)

    inpainted_output = self.inpainter(unet_output , binary_mask)

    return inpainted_output
</code></pre>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/aniruth_sundararaja1">@Aniruth_Sundararaja1</a>,</p>
<p>When you re-wrap the output tensor from your <code>Unet</code>, you destroy the history of operations and hence don’t have a gradient. Make sure to not move your tensor from <code>torch</code> to <code>numpy</code> (and back again to <code>torch</code>).</p>
<aside class="quote no-group" data-username="Aniruth_Sundararaja1" data-post="1" data-topic="203275">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aniruth_sundararaja1/48/61218_2.png" class="avatar"> Aniruth_Sundararaja1:</div>
<blockquote>
<pre><code class="lang-auto">    difference = (torch.from_numpy(flary_img_gray) - torch.from_numpy(unet_output_gray))
</code></pre>
</blockquote>
</aside>
<p>Also, you should just be able to use <code>binary_mask = torch.where(difference_tensor &gt; self.learnable_threshold, 1, 0)</code> (without creating new tensor), as scalars can easily broadcast across your multi-dimensional tensor.</p>
<aside class="quote no-group" data-username="Aniruth_Sundararaja1" data-post="1" data-topic="203275">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aniruth_sundararaja1/48/61218_2.png" class="avatar"> Aniruth_Sundararaja1:</div>
<blockquote>
<pre><code class="lang-auto">    binary_mask = torch.where(difference_tensor &gt; self.learnable_threshold, torch.tensor(1.0).to(options.device), torch.tensor(0.0).to(options.device))
</code></pre>
</blockquote>
</aside>

Aniruth_Sundararaja1: <p>Thanks <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a></p>
<p>I tried out the solution as you said  , but my param is still not changing<br>
I am attaching my updated code as you has told for reference</p>
<p>def tensor_to_cv2_gray(tensor):<br>
tensor = tensor.detach().cpu()<br>
tensor = tensor.permute(0, 2, 3, 1)</p>
<pre><code>r_factor = torch.tensor(0.2989, dtype=tensor.dtype, device=tensor.device)
g_factor = torch.tensor(0.5870, dtype=tensor.dtype, device=tensor.device)
b_factor = torch.tensor(0.1140, dtype=tensor.dtype, device=tensor.device)

gray_tensor = tensor[..., 0] * r_factor + tensor[..., 1] * g_factor + tensor[..., 2] * b_factor

gray_tensor = gray_tensor.unsqueeze(1).permute(0, 3, 1, 2)

return gray_tensor
</code></pre>
<h1><a name="dataloader-1" class="anchor" href="#dataloader-1"></a>DataLoader</h1>
<p>train_dataset, train_dataloader = training_dataloader(options.batch_size, options.no_of_workers, options.train_dataset_path, options.image_size)</p>
<h1><a name="models-2" class="anchor" href="#models-2"></a>Models</h1>
<p>def load_pretrained_weights(original_model, new_model):<br>
orig_ = {}<br>
new_ = {}<br>
for (name, param) in original_model.named_parameters():<br>
orig_[name] = param<br>
for (name, param) in new_model.named_parameters():<br>
if name in orig_:<br>
new_[name] = orig_[name]<br>
else:<br>
new_[name] = torch.randn_like(param) / torch.norm(torch.randn_like(param))<br>
return new_</p>
<p>class Combined_Model(nn.Module):</p>
<pre><code>def __init__(self , options):
    super(Combined_Model, self).__init__()

    self.pretrained_state_dict = torch.load(os.path.join(options.pretrained, 'G0000000.pt'), map_location=torch.device('cuda'))
    self.unet = UNet().to(options.device)


    if options.with_prompts:
        self.inpainter = Prompted_InpaintGenerator(options)
        self.org_gan = InpaintGenerator(options)
        #self.inpainter.load_state_dict(load_pretrained_weights(self.org_gan, self.pretrained_state_dict), strict=False)
        self.inpainter.load_state_dict(load_pretrained_weights(self.org_gan , self.inpainter) , strict=True)
    else:
        self.inpainter = InpaintGenerator(options)
        self.inpainter.load_state_dict(torch.load(os.path.join(options.pretrained, 'G0000000.pt'), map_location=options.device), strict=False)

    self.models = [self.unet, self.inpainter]

    #self.learnable_threshold = nn.Parameter(torch.tensor(0.5), requires_grad=True)
    self.learnable_threshold = nn.Parameter(torch.tensor(0.5), requires_grad=True)



def forward(self , x):

    unet_output = self.unet(x)
    unet_output_gray = tensor_to_cv2_gray(unet_output)
    flary_img_gray = tensor_to_cv2_gray(x)
    #    difference = (torch.from_numpy(flary_img_gray) - torch.from_numpy(unet_output_gray))
    difference = (flary_img_gray) - (unet_output_gray)
    #difference_tensor = torch.tensor(difference, dtype=torch.float32).to(options.device)
    difference_tensor = difference.clone().to(options.device)
    binary_mask = torch.where(difference_tensor &gt; self.learnable_threshold, 1 , 0)
    #binary_mask = self.learnable_threshold_layer(difference_tensor, options.device)
    #binary_mask = binary_mask.unsqueeze(1)
    binary_mask = binary_mask.permute(0 , 2 , 1 , 3)
    inpainted_output = self.inpainter(unet_output , binary_mask)

    return inpainted_output
</code></pre>
<p>For your reference , this is the output i get<br>
0%|                                                                                                              | 0/100 [00:00&lt;?, ?it/s]torch.Size([2, 3, 256, 256]) torch.Size([2, 1, 256, 256])<br>
Epoch Loss : 0.00016948020493146032<br>
Learnable Threshold: 0.5<br>
1%|▉                                                                                                  | 1/100 [02:40&lt;4:25:12, 160.73s/it]torch.Size([2, 3, 256, 256]) torch.Size([2, 1, 256, 256])<br>
Epoch Loss : 0.000250783225055784<br>
Learnable Threshold: 0.5</p>
<p>Could you please look and tell , where i am going wrong</p>

Aniruth_Sundararaja1: <p>Update</p>
<p>when i tried checking in model names parameters , i could not find the learnble threshold parameter</p>