Hello!
For the last 2 weeks I’ve been stuck trying to count balls from a synthetic dataset I generated.
When I set a batch size higher than 1, the network predicts the average value all the time.
Otherwise, the network works great in train/val

Dataset: 5000 images with grey background and blue balls of the same size around. There are between 1 and 10 balls per image.
Optimizer: Adam with learning rate of 1e-4
Normalization: for the normalization of the images I used:

transforms.Compose([
        transforms.Resize((128,128)),
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5],
                             [0.5, 0.5, 0.5])
    ])


Loss: L1 (Mean absolute error)
Network:

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(3, 6, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2))
        self.layer2 = nn.Sequential(
            nn.Conv2d(6, 15, 5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2))
        self.dropout = nn.Dropout(0.2)
        self.fc1 = nn.Linear(12615, 1000)
        self.fc2 = nn.Linear(1000, 1)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.reshape(out.size(0), -1)
        out = self.dropout(out)
        out = self.fc1(out)
        out = self.dropout(out)
        out = self.fc2(out)
        return out

I am bit frustrated because this is the easiest example of object counting and I can not make it work properly.
How can I avoid predicting the average value (5) when raising the batch size?
PD: I tried many things: cyclic LR, Batchnorm2d, deeper networks, all the possible optimizers…
Thanks!