PyTorch Inequality Gradient

Viewed 154

I wrote a PyTorch model roughly as follows:

import torch.nn as nn
class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.layer1 = nn.Sequential(nn.Linear(64 * 64, 16), nn.LeakyReLU(0.2))
        self.layer2 = nn.Sequential(nn.Linear(16, 32), nn.LeakyReLU(0.2))
        self.layer3 = nn.Sequential(nn.Linear(32, 64), nn.LeakyReLU(0.2))
        self.layer4 = nn.Sequential(nn.Linear(64, 15), nn.Tanh())
    def forward(self, x):
        return (self.layer4(self.layer3(self.layer2(self.layer1(x)))) < 0).float()

Notice what I want to do: I want forward to return a tensor of 0s and 1s. However, this does not train, probably because the derivative of the inequality is zero.

How can I make a model like this one train, for example, if I want to do image segmentation?

1 Answers

As you said, you can't train something like x<0.
You should be fine even if you get rid of the <0 part and use

    return self.layer4(self.layer3(self.layer2(self.layer1(x))))

as long as you are using the appropriate loss. I think what you would want to use is nn.BCEWithLogitsLoss. In that case you should get the Tanh out of the last layer since nn.BCEWithLogitsLoss internally computes with sigmoid.
(There are options of using nn.BCEloss() with sigmoid at the last layer, or even stick with Tanh, but I don't think there's a reason to take the long way.)

So in the training phase, the neural network tries as hard as it could to fit the output to 0s and 1s. After that, it is the testing phase that you should take the output of the layer, and give it some kind of threshold to change the values to precisely 1s and 0s.(like you did (output<0).float())

You will find useful sources if you search for multilabel classification.

Related