spanskiduh
/
uav-localization
1
1
Fork 0
Code Issues Pull Requests Packages Projects Releases 1 Wiki Activity

Add dropout layers

main
Gašper Spagnolo 2023-03-19 19:57:16 +01:00
parent bcbbea6cd2
commit 3501bea445
1 changed files with 12 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
code/autoencoder.py
View File

@ -130,7 +130,12 @@ class Encoder(nn.Module):
super().__init__()
self.debug = debug
self.net = nn.Sequential(
self.linear = nn.Sequential(
nn.Flatten(),
nn.Linear(IMG_H * IMG_W, latent_dim),
)
self.conv = nn.Sequential(
nn.Conv2d(
in_channels=in_channels,
out_channels=out_channels,
@ -138,6 +143,7 @@ class Encoder(nn.Module):
stride=stride,
),
nn.BatchNorm2d(out_channels),
nn.Dropout(0.3),
act_fn,
nn.Conv2d(
in_channels=out_channels,
@ -146,6 +152,7 @@ class Encoder(nn.Module):
stride=stride,
),
nn.BatchNorm2d(out_channels * 2),
nn.Dropout(0.2),
act_fn,
nn.Conv2d(
in_channels=out_channels * 2,
@ -154,6 +161,7 @@ class Encoder(nn.Module):
stride=stride,
),
nn.BatchNorm2d(out_channels * 4),
nn.Dropout(0.1),
act_fn,
nn.Conv2d(
in_channels=out_channels * 4,
@ -171,8 +179,6 @@ class Encoder(nn.Module):
),
act_fn,
nn.BatchNorm2d(out_channels * 8),
nn.Flatten(),
nn.Linear(IMG_H * IMG_W, latent_dim),
)
def forward(self, x):
@ -182,7 +188,8 @@ class Encoder(nn.Module):
# x = layer(x)
# if self.debug:
# print(layer.__class__.__name__, "output shape:\t", x.shape)
encoded_latent_image = self.net(x)
encoded_latent_image = self.conv(x)
encoded_latent_image = self.linear(encoded_latent_image)
return encoded_latent_image
@ -247,7 +254,7 @@ class Decoder(nn.Module):
kernel_size=kernel_size,
stride=stride,
),
act_fn,
nn.ReLU(),
nn.BatchNorm2d(in_channels),
)
@ -323,7 +330,6 @@ class ConvolutionalAutoencoder:
for epoch in range(epochs):
print(f"Epoch {epoch+1}/{epochs}")
train_losses = []
# ------------
# TRAINING