Simpler convolutions, aligned to the mouse grid
@@ -10,11 +10,10 @@ class Model(nn.Module):
super(Model, self).__init__()
self.conv = nn.Sequential(
- nn.Conv2d(input_shape[0], 32, 4, stride=3),
+ nn.Conv2d(input_shape[0], 64, 5, stride=5),
nn.ReLU(),
- nn.Conv2d(32, 64, 4, stride=2),
+ nn.Conv2d(64, 64, 2, stride=1),
nn.ReLU(),
- nn.Conv2d(64, 64, 3, stride=2),
)
conv_out_size = self._get_conv_out(input_shape)
