Hello,
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 = inputData.to(device)
            optimizer.zero_grad()
            outputData ,a = net(inputData)
            loss_MSE = criterion(outputData, inputData)
            loss_TVPenalty = TV_Penalty(net.DEC)
            loss = loss_MSE + loss_TVPenalty
            loss.backward()
            optimizer.step()
            net.DEC.weight.data.clamp_(min=0)
            running_loss += loss.item()
        loss = running_loss / len(trainloader)
        train_loss.append(loss)
    return train_loss

def TV_Penalty(Layer):
    loss = torch.pow((Layer.weight.data[1:]-Layer.weight.data[:-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.