Intermediate output should be in [-1, 1] but tanh activation saturates - solution?

Viewed 120

I am prediction a set of coordinates on a grid at an intermediate layer of my neural network. This coordinates should be in range [-1, 1]. My idea was to use a tanh activation function to achieve that - but unfortunately, by doing so after a few steps the output is always -1 or 1. I read, that this can be prevented by using a batch normalization before the tanh. This does indeed work, but at a cost - it forces all predicted coordinates towards the centre, which is a huge and a wrong bias, as I am predicting body landmarks, which are not necessarily in the middle of the image.

Does anyone have an idea, what I can do? The way I predict the coordinates is the following:

class Encoder(nn.Module):
    def __init__(self, n_parts, n_features, residual_dim, reconstr_dim, depth_s, depth_a, device, p_dropout=0.2, c=16, dim=64):
        super(Encoder, self).__init__()
        self.preprocessor = PreProcessor(residual_dim, reconstr_dim)
        self.k = n_parts
        self.dropout = nn.Dropout(p_dropout)
        self.tanh = nn.Tanh()
        self.relu = nn.ReLU()
        self.relu = nn.LeakyReLU()
        self.bn1 = nn.BatchNorm2d(c * self.k)
        self.bn2 = nn.BatchNorm2d(2 * n_parts)
        self.hg_shape = Hourglass(depth_s, residual_dim)
        self.to_channels = Conv(residual_dim, c * n_parts, kernel_size=3, stride=1, bn=False, relu=False)
        self.mu = nn.Conv2d(in_channels=c * n_parts, out_channels=2 * n_parts, kernel_size=dim, groups=n_parts)

    def forward(self, img):
        bn = img.shape[0]
        img_preprocessed = self.preprocessor(img) # shape [bn, 256, 64, 64]

        # Get Shape Representation
        img_shape = self.hg_shape(img_preprocessed)
        img_shape = self.dropout(img_shape)
        img_c = self.to_channels(img_shape)
        img_c = self.relu(self.bn1(img_c))
        mu = self.mu(img_c)
        mu = self.tanh(mu).reshape(bn, self.k, 2)
        
        return mu
0 Answers
Related