TL;DR
I am unable to overfit batches with multiple samples using autoencoder.
Fully connected decoder seems to handle more samples per batch than conv decoder, but then also fails when number of samples increases.
Why is this happening, and how to debug this?
In depth
I am trying to use an auto encoder on 1d data points of size (n, 1, 1024), where n is the number of samples in the batch.
I am trying to overfit to that single batch.
Using a convolutional decoder, I am only able to fit a single sample (n=1), and when n>1 I am unable to drop the loss (MSE) below 0.2.
In blue: expected output (=input), in orange: reconstruction.
Single sample, single batch:
Multiple samples, single batch, loss won't go down:
Using more than one sample, we can see the net learns the general shape of the input (=output) signal, but greatly misses by an over-all constant.
Using a fully connected decoder does manage to reconstruct batches of multiple samples:
Relevant code:
class Conv1DBlock(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size):
super().__init__()
self._in_channels = in_channels
self._out_channels = out_channels
self._kernel_size = kernel_size
self._block = nn.Sequential(
nn.Conv1d(
in_channels=self._in_channels,
out_channels=self._out_channels,
kernel_size=self._kernel_size,
stride=1,
padding=(self._kernel_size - 1) // 2,
),
# nn.BatchNorm1d(num_features=out_channels),
nn.ReLU(True),
nn.MaxPool1d(kernel_size=2, stride=2),
)
def forward(self, x):
for layer in self._block:
x = layer(x)
return x
class Upsample1DBlock(nn.Module):
def __init__(self, in_channels, out_channels, factor):
super().__init__()
self._in_channels = in_channels
self._out_channels = out_channels
self._factor = factor
self._block = nn.Sequential(
nn.Conv1d(
in_channels=self._in_channels,
out_channels=self._out_channels,
kernel_size=3,
stride=1,
padding=1
), # 'same'
nn.ReLU(True),
nn.Upsample(scale_factor=self._factor, mode='linear', align_corners=True),
)
def forward(self, x):
x_tag = x
for layer in self._block:
x_tag = layer(x_tag)
# interpolated = F.interpolate(x, scale_factor=0.5, mode='linear') # resnet idea
return x_tag
encoder:
self._encoder = nn.Sequential(
# n, 1024
nn.Unflatten(dim=1, unflattened_size=(1, 1024)),
# n, 1, 1024
Conv1DBlock(in_channels=1, out_channels=8, kernel_size=15),
# n, 8, 512
Conv1DBlock(in_channels=8, out_channels=16, kernel_size=11),
# n, 16, 256
Conv1DBlock(in_channels=16, out_channels=32, kernel_size=7),
# n, 32, 128
Conv1DBlock(in_channels=32, out_channels=64, kernel_size=5),
# n, 64, 64
Conv1DBlock(in_channels=64, out_channels=128, kernel_size=3),
# n, 128, 32
nn.Conv1d(in_channels=128, out_channels=128, kernel_size=32, stride=1, padding=0), # FC
# n, 128, 1
nn.Flatten(start_dim=1, end_dim=-1),
# n, 128
)
conv decoder:
self._decoder = nn.Sequential(
nn.Unflatten(dim=1, unflattened_size=(128, 1)), # 1
Upsample1DBlock(in_channels=128, out_channels=64, factor=4), # 4
Upsample1DBlock(in_channels=64, out_channels=32, factor=4), # 16
Upsample1DBlock(in_channels=32, out_channels=16, factor=4), # 64
Upsample1DBlock(in_channels=16, out_channels=8, factor=4), # 256
Upsample1DBlock(in_channels=8, out_channels=1, factor=4), # 1024
nn.ReLU(True),
nn.Conv1d(in_channels=1, out_channels=1, kernel_size=3, stride=1, padding=1),
nn.ReLU(True),
nn.Flatten(start_dim=1, end_dim=-1),
nn.Linear(1024, 1024)
)
FC decoder:
self._decoder = nn.Sequential(
nn.Linear(128, 256),
nn.ReLU(True),
nn.Linear(256, 512),
nn.ReLU(True),
nn.Linear(512, 1024),
nn.ReLU(True),
nn.Flatten(start_dim=1, end_dim=-1),
nn.Linear(1024, 1024)
)
Another observation is that when the batch size increases more, to say, 16, the FC decoder also starts to fail.
In the image, 4 samples of a 16 sample batch I am trying to overfit
What could be wrong with the conv decoder?
How to debug this or make the conv decoder work?
Please notice that the same infrastructure with only the encoder and decoder different do manage to overfit and generalize over MNIST.
(This is also posted here, but I think this is still ok to do. If not, please tell me and I will delete one).
