Layer weights in custom Loss results in autograd problem


I want to introduce a kind of total variation penalty.

I want to realize this by the fact, that the neural network is penalized for it, if the neighboured weights differ strongly.

But I want to apply this only in one layer. So I have added this additionally in the loss function.

My code parts are therefore as follows:

class CNN(nn.Module):
   def __init__(self):
       super(CNN, self).__init__()

       # decoder 
       self.conv1 = nn.Conv1d(in_channels=1, out_channels=1, kernel_size=50, bias=False)
       self.conv2 = nn.BatchNorm1d(1)
       self.conv3 = nn.Conv1d(in_channels=1,out_channels=1, kernel_size=30, bias=False)
       self.conv4 = nn.BatchNorm1d(1)
       self.conv5 = nn.Conv1d(in_channels=1,out_channels=1, kernel_size=11, bias=False)
       self.conv6 = nn.BatchNorm1d(1)
       self.ASC = nn.Softmax(dim=1)
       # encoder 
       self.DEC = nn.Linear(in_features=3,out_features=91)

   def forward(self, x):
       x = F.leaky_relu(self.conv1(x))
       x = F.leaky_relu(self.conv2(x))
       x = F.leaky_relu(self.conv3(x))
       x = F.leaky_relu(self.conv4(x))
       x = F.leaky_relu(self.conv5(x))
       x = x*3.5
       x = F.leaky_relu(self.conv6(x))
       a = self.ASC(x)
       x = F.leaky_relu(self.DEC(a))
       return x, a
net = CNN()
def train(net, trainloader, NUM_EPOCHS, cur_epoch=0):
    train_loss = []
    for epoch in range(cur_epoch, NUM_EPOCHS):
        running_loss = 0.0
        for data in trainloader:
            inputData = data
            inputData =
            outputData ,a = net(inputData)
            loss_MSE = criterion(outputData, inputData)
            loss_TVPenalty = TV_Penalty(net.DEC)
            loss = loss_MSE + loss_TVPenalty
            running_loss += loss.item()
        loss = running_loss / len(trainloader)
    return train_loss
def TV_Penalty(Layer):
    loss = torch.pow(([1:][:-1]), 2).sum()*TV_PENALTY_CONSTANT
    return loss

If I use it this way, the autograde function of Pytorch is not applied to the python def TV_Penalty(Layer): function.

How to use the autograd in Pytocht on the structure of the weights of the network itself? So then this is considered in an additional loss function.

You should be able to use the .weight parameters directly, instead of their .data attribute, which Autograd cannot track and which should thus be avoided. :wink:

Hihih, many thanks, works like a charm. :grinning: