I want the loss function like ||Mask · (output − target )||2. Should the below code works fine for gradient flow? Note: Mask is bool(True or False).
import torch.nn as nn
import torch.nn.functional as F
import torch
from torchsummary import summary
class TwoLayerNet(nn.Module):
def __init__(self):
super(TwoLayerNet, self).__init__()
conv_layers=[]
self.conv=[]
conv_layers.append(nn.Sequential(nn.Conv2d(3, 5, kernel_size=3, stride=1,padding=1, bias=True) ) )
conv_layers.append(nn.Sequential(nn.Conv2d(5, 5, kernel_size=3, stride=1,padding=1, bias= True),nn.LeakyReLU(0.2,True) ))
self.conv = nn.Sequential(*conv_layers)
def forward(self, x):
x1=self.conv(x)
return x1
torch.manual_seed(88)
x = torch.ones(1,3,2,2)
tar=torch.ones(1,5,2,2)*0.1
model = TwoLayerNet()
y_pred = model(x)
mask=y_pred>tar
print("mask",mask)
crit=torch.nn.MSELoss()
loss= mask*crit(y_pred,tar)