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?