Skip to content

Commit 2f1b321

Browse files
renan-siqueirarenan-siqueira
committed
feat: create more models of autoencoder (convAE, VAE, convVAE)
1 parent 84cd3e1 commit 2f1b321

File tree

3 files changed

+208
-33
lines changed

3 files changed

+208
-33
lines changed

‎models/autoencoder.py‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
1+
import torch
12
import torch.nn as nn
3+
import torch.nn.functional as F
24

35

6+
# Autoencoder Linear
47
class Autoencoder(nn.Module):
58
def __init__(self, input_dim, encoding_dim):
69
super(Autoencoder, self).__init__()
@@ -24,3 +27,127 @@ def forward(self, x):
2427
x = self.encoder(x)
2528
x = self.decoder(x)
2629
return x
30+
31+
32+
# Autoencoder Convolucional
33+
class ConvolutionalAutoencoder(nn.Module):
34+
def __init__(self):
35+
super(ConvolutionalAutoencoder, self).__init__()
36+
37+
# Encoder
38+
self.enc1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
39+
self.enc2 = nn.Conv2d(64, 32, kernel_size=3, padding=1)
40+
self.enc3 = nn.Conv2d(32, 16, kernel_size=3, padding=1)
41+
self.pool = nn.MaxPool2d(2, 2, return_indices=True)
42+
43+
# Decoder
44+
self.dec1 = nn.ConvTranspose2d(16, 32, kernel_size=2, stride=2)
45+
self.dec2 = nn.ConvTranspose2d(32, 64, kernel_size=2, stride=2)
46+
self.dec3 = nn.ConvTranspose2d(64, 3, kernel_size=2, stride=2)
47+
48+
def forward(self, x):
49+
x, idxs1 = self.pool(F.relu(self.enc1(x)))
50+
x, idxs2 = self.pool(F.relu(self.enc2(x)))
51+
x, idxs3 = self.pool(F.relu(self.enc3(x)))
52+
53+
x = F.relu(self.dec1(x))
54+
x = F.relu(self.dec2(x))
55+
x = torch.sigmoid(self.dec3(x))
56+
return x
57+
58+
59+
# Variational Autoencoder
60+
class VariationalAutoencoder(nn.Module):
61+
def __init__(self, encoding_dim=128):
62+
super(VariationalAutoencoder, self).__init__()
63+
64+
# Encoder
65+
self.enc1 = nn.Linear(3 * 64 * 64, 512)
66+
self.enc2 = nn.Linear(512, 256)
67+
self.enc3 = nn.Linear(256, encoding_dim)
68+
69+
# Latent space
70+
self.fc_mu = nn.Linear(encoding_dim, encoding_dim)
71+
self.fc_log_var = nn.Linear(encoding_dim, encoding_dim)
72+
73+
# Decoder
74+
self.dec1 = nn.Linear(encoding_dim, encoding_dim)
75+
self.dec2 = nn.Linear(encoding_dim, 256)
76+
self.dec3 = nn.Linear(256, 512)
77+
self.dec4 = nn.Linear(512, 3 * 64 * 64)
78+
79+
def reparameterize(self, mu, log_var):
80+
std = torch.exp(0.5 * log_var)
81+
eps = torch.randn_like(std)
82+
return mu + eps * std
83+
84+
def forward(self, x):
85+
x = F.relu(self.enc1(x))
86+
x = F.relu(self.enc2(x))
87+
x = F.relu(self.enc3(x))
88+
89+
mu = self.fc_mu(x)
90+
log_var = self.fc_log_var(x)
91+
z = self.reparameterize(mu, log_var)
92+
93+
x = F.relu(self.dec1(z))
94+
x = F.relu(self.dec2(x))
95+
x = F.relu(self.dec3(x))
96+
x = torch.sigmoid(self.dec4(x))
97+
98+
return x, mu, log_var
99+
100+
101+
# Convolucional Variational Autoencoder
102+
class ConvolutionalVAE(nn.Module):
103+
def __init__(self):
104+
super(ConvolutionalVAE, self).__init__()
105+
106+
# Encoder
107+
self.enc1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
108+
self.enc2 = nn.Conv2d(64, 32, kernel_size=3, padding=1)
109+
self.enc3 = nn.Conv2d(32, 16, kernel_size=3, padding=1)
110+
self.pool = nn.MaxPool2d(2, 2)
111+
112+
self.fc_mu = nn.Linear(16 * 8 * 8, 128)
113+
self.fc_log_var = nn.Linear(16 * 8 * 8, 128)
114+
115+
# Decoder
116+
self.decoder_input = nn.Linear(128, 16 * 8 * 8)
117+
self.dec1 = nn.ConvTranspose2d(16, 32, kernel_size=3, padding=1)
118+
self.dec2 = nn.ConvTranspose2d(32, 64, kernel_size=3, padding=1)
119+
self.dec3 = nn.ConvTranspose2d(64, 3, kernel_size=3, padding=1)
120+
121+
self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
122+
123+
def reparameterize(self, mu, log_var):
124+
std = torch.exp(0.5 * log_var)
125+
eps = torch.randn_like(std)
126+
return mu + eps * std
127+
128+
def forward(self, x):
129+
# Encoding
130+
x = F.relu(self.enc1(x))
131+
x = self.pool(x)
132+
x = F.relu(self.enc2(x))
133+
x = self.pool(x)
134+
x = F.relu(self.enc3(x))
135+
x = self.pool(x)
136+
137+
x = x.view(x.size(0), -1) # Flatten
138+
139+
mu = self.fc_mu(x)
140+
log_var = self.fc_log_var(x)
141+
z = self.reparameterize(mu, log_var)
142+
143+
# Decoding
144+
x = self.decoder_input(z)
145+
x = x.view(x.size(0), 16, 8, 8) # Unflatten
146+
x = self.upsample(x)
147+
x = F.relu(self.dec1(x))
148+
x = self.upsample(x)
149+
x = F.relu(self.dec2(x))
150+
x = self.upsample(x)
151+
x = torch.sigmoid(self.dec3(x))
152+
153+
return x, mu, log_var

‎run.py‎

Lines changed: 33 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
11
import os
22
import torch
33

4-
from models.autoencoder import Autoencoder
4+
from models.autoencoder import Autoencoder, ConvolutionalAutoencoder, VariationalAutoencoder, ConvolutionalVAE
55
from utils.dataloader import get_dataloader
6-
from utils.trainer import train_autoencoder, visualize_reconstructions, save_checkpoint, load_checkpoint, evaluate_autoencoder
6+
from utils.trainer import train_autoencoder, visualize_reconstructions, load_checkpoint, evaluate_autoencoder
77
from settings import settings
88

99

10-
def main(load_trained_model):
10+
def main(load_trained_model, ae_type='ae'):
1111
BATCH_SIZE = 32
1212
INPUT_DIM = 3 * 64 * 64
1313
ENCODING_DIM = 64
@@ -16,25 +16,46 @@ def main(load_trained_model):
1616
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
1717
dataloader = get_dataloader(settings.DATA_PATH, BATCH_SIZE)
1818

19-
model = Autoencoder(INPUT_DIM, ENCODING_DIM).to(device)
19+
if ae_type == 'ae':
20+
model = Autoencoder(INPUT_DIM, ENCODING_DIM).to(device)
21+
elif ae_type == 'conv':
22+
model = ConvolutionalAutoencoder().to(device)
23+
elif ae_type == 'vae':
24+
model = VariationalAutoencoder().to(device)
25+
elif ae_type == 'conv_vae':
26+
model = ConvolutionalVAE().to(device)
27+
else:
28+
raise ValueError(f"Unknown AE type: {ae_type}")
29+
2030
optimizer = torch.optim.Adam(model.parameters())
2131

2232
start_epoch = 0
2333
if os.path.exists(settings.PATH_SAVED_MODEL):
24-
model, optimizer, start_epoch = load_checkpoint(model, optimizer, settings.PATH_SAVED_MODEL, device)
34+
model, optimizer, start_epoch = load_checkpoint(
35+
model, optimizer, settings.PATH_SAVED_MODEL, device
36+
)
2537
print(f"Loaded checkpoint and continuing training from epoch {start_epoch}.")
2638

2739
if not load_trained_model:
28-
for epoch in range(start_epoch, NUM_EPOCHS):
29-
train_autoencoder(model, dataloader, device=device)
30-
print(f"Epoch [{epoch + 1}/{NUM_EPOCHS}] complete!")
31-
save_checkpoint(model, optimizer, epoch, settings.PATH_SAVED_MODEL)
40+
train_autoencoder(
41+
model,
42+
dataloader,
43+
num_epochs=NUM_EPOCHS,
44+
device=device,
45+
start_epoch=start_epoch,
46+
optimizer=optimizer,
47+
ae_type=ae_type
48+
)
49+
print(f"Training complete up to epoch {NUM_EPOCHS}!")
3250

3351
valid_dataloader = get_dataloader(settings.VALID_DATA_PATH, BATCH_SIZE)
34-
avg_valid_loss = evaluate_autoencoder(model, valid_dataloader, device)
52+
avg_valid_loss = evaluate_autoencoder(model, valid_dataloader, device, ae_type)
3553
print(f"Average validation loss: {avg_valid_loss:.4f}")
36-
visualize_reconstructions(model, valid_dataloader, num_samples=10, device=device)
54+
55+
visualize_reconstructions(
56+
model, valid_dataloader, num_samples=10, device=device, ae_type=ae_type
57+
)
3758

3859

3960
if __name__ == "__main__":
40-
main(False)
61+
main(False, ae_type='conv_vae')

‎utils/trainer.py‎

Lines changed: 48 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -2,21 +2,29 @@
22
import torch
33
import torch.optim as optim
44
import torch.nn as nn
5+
import torch.nn.functional as F
56
from torchvision.utils import save_image, make_grid
67
import matplotlib.pyplot as plt
78

89

9-
def train_autoencoder(model, dataloader, num_epochs=5, learning_rate=0.001, device='cpu', start_epoch=0, optimizer=None):
10+
def train_autoencoder(model, dataloader, num_epochs=5, learning_rate=0.001, device='cpu', start_epoch=0, optimizer=None, ae_type='ae'):
1011
criterion = nn.MSELoss()
1112
if optimizer is None:
1213
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
1314

1415
for epoch in range(start_epoch, num_epochs):
1516
for data in dataloader:
1617
img = data.to(device)
17-
img = img.view(img.size(0), -1)
18-
output = model(img)
19-
loss = criterion(output, img)
18+
19+
if ae_type not in ['conv', 'conv_vae']:
20+
img = img.view(img.size(0), -1)
21+
22+
if ae_type in ['vae', 'conv_vae']:
23+
recon_x, mu, log_var = model(img)
24+
loss = loss_function_vae(recon_x, img, mu, log_var)
25+
else:
26+
output = model(img)
27+
loss = criterion(output, img)
2028

2129
optimizer.zero_grad()
2230
loss.backward()
@@ -28,12 +36,45 @@ def train_autoencoder(model, dataloader, num_epochs=5, learning_rate=0.001, devi
2836
return model
2937

3038

31-
def visualize_reconstructions(model, dataloader, num_samples=10, device='cpu', save_path="./samples"):
39+
def loss_function_vae(recon_x, x, mu, log_var):
40+
BCE = F.binary_cross_entropy(recon_x, x, reduction='sum')
41+
KLD = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
42+
return BCE + KLD
43+
44+
45+
def evaluate_autoencoder(model, dataloader, device, ae_type):
46+
model.eval()
47+
total_loss = 0
48+
criterion = nn.MSELoss()
49+
with torch.no_grad():
50+
for data in dataloader:
51+
img = data.to(device)
52+
53+
if ae_type not in ['conv', 'conv_vae']:
54+
img = img.view(img.size(0), -1)
55+
56+
if ae_type in ['vae', 'conv_vae']:
57+
output, _, _ = model(img)
58+
else:
59+
output = model(img)
60+
loss = criterion(output, img)
61+
total_loss += loss.item()
62+
63+
return total_loss / len(dataloader)
64+
65+
66+
def visualize_reconstructions(model, dataloader, num_samples=10, device='cpu', save_path="./samples", ae_type='ae'):
3267
model.eval()
3368
samples = next(iter(dataloader))
3469
samples = samples[:num_samples].to(device)
35-
samples = samples.view(samples.size(0), -1)
36-
reconstructions = model(samples)
70+
71+
if ae_type not in ['conv', 'conv_vae']:
72+
samples = samples.view(samples.size(0), -1)
73+
74+
if ae_type in ['vae', 'conv_vae']:
75+
reconstructions, _, _ = model(samples)
76+
else:
77+
reconstructions = model(samples)
3778

3879
samples = samples.view(-1, 3, 64, 64)
3980
reconstructions = reconstructions.view(-1, 3, 64, 64)
@@ -75,17 +116,3 @@ def load_checkpoint(model, optimizer, path, device):
75116
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
76117
epoch = checkpoint['epoch']
77118
return model, optimizer, epoch + 1
78-
79-
80-
def evaluate_autoencoder(model, dataloader, device):
81-
model.eval()
82-
total_loss = 0
83-
criterion = nn.MSELoss()
84-
with torch.no_grad():
85-
for data in dataloader:
86-
img = data.to(device)
87-
img = img.view(img.size(0), -1)
88-
output = model(img)
89-
loss = criterion(output, img)
90-
total_loss += loss.item()
91-
return total_loss / len(dataloader)

0 commit comments

Comments
 (0)