132 lines
4.9 KiB
Python
132 lines
4.9 KiB
Python
import numpy as np
|
|
import torch
|
|
from torch.utils import data
|
|
from typing import Tuple
|
|
import matplotlib.pyplot as plt
|
|
from stack_with_padding import stack_with_padding
|
|
from torch.backends import mps
|
|
from tqdm import tqdm
|
|
|
|
|
|
def training_loop(
|
|
network: torch.nn.Module,
|
|
train_data: torch.utils.data.Dataset,
|
|
eval_data: torch.utils.data.Dataset,
|
|
num_epochs: int,
|
|
show_progress: bool = False
|
|
) -> Tuple[list, list]:
|
|
optimizer = torch.optim.Adam(params=network.parameters(), lr=1e-5, weight_decay=1e-5)
|
|
data_loader_train = torch.utils.data.DataLoader(
|
|
dataset=train_data,
|
|
batch_size=64,
|
|
shuffle=True,
|
|
collate_fn=stack_with_padding,
|
|
num_workers=10,
|
|
pin_memory=True
|
|
)
|
|
data_loader_eval = torch.utils.data.DataLoader(
|
|
dataset=eval_data,
|
|
batch_size=64,
|
|
collate_fn=stack_with_padding,
|
|
num_workers=10,
|
|
pin_memory=True
|
|
)
|
|
# cuda can be added if executed on machine with cuda available
|
|
use_mps = torch.backends.mps.is_available()
|
|
device = torch.device("mps" if use_mps else "cpu")
|
|
network.to(device)
|
|
epoch_losses = []
|
|
eval_losses = []
|
|
# init progress bar
|
|
progress_bar = None
|
|
if show_progress:
|
|
total_iterations = num_epochs * (len(data_loader_train) + len(data_loader_eval))
|
|
progress_bar = tqdm(total=total_iterations, desc="Training and Evaluation")
|
|
for epoch in range(num_epochs):
|
|
mbl = []
|
|
network.train()
|
|
for inputs, known, targets, _ in data_loader_train:
|
|
inputs = inputs.to(device).float()
|
|
known = known.to(device).float()
|
|
targets = targets.to(device).float()
|
|
optimizer.zero_grad()
|
|
output = network(inputs)
|
|
known_tensor = known.bool()
|
|
# Extract only pixelated area from output / targets for loss calculation
|
|
target = targets[known_tensor]
|
|
out = output[known_tensor]
|
|
rmse_loss = torch.nn.MSELoss()(out, target)
|
|
rmse_loss = torch.sqrt(rmse_loss)
|
|
rmse_loss.backward()
|
|
optimizer.step()
|
|
|
|
mbl.append(rmse_loss.item())
|
|
if show_progress:
|
|
progress_bar.update(1)
|
|
epoch_losses.append(np.mean(mbl))
|
|
# Save model state after each epoch, can be optimized, good for now
|
|
torch.save(network.state_dict(), f"model_{epoch + 1}.pth")
|
|
|
|
network.eval()
|
|
with torch.no_grad():
|
|
eval_mlb = []
|
|
for inputs, known, targets, _ in data_loader_eval:
|
|
inputs = inputs.to(device).float()
|
|
known = known.to(device).float()
|
|
targets = targets.to(device).float()
|
|
output = network(inputs)
|
|
known_tensor = known.bool()
|
|
target = targets[known_tensor]
|
|
out = output[known_tensor]
|
|
loss = torch.nn.MSELoss()(out, target)
|
|
loss = torch.sqrt(loss)
|
|
eval_mlb.append(loss.item())
|
|
if show_progress:
|
|
progress_bar.update(1)
|
|
eval_losses.append(np.mean(eval_mlb))
|
|
return epoch_losses, eval_losses
|
|
|
|
|
|
if __name__ == "__main__":
|
|
from model import SimpleCNN
|
|
from create_dataset import RandomImagePixelationDataset
|
|
|
|
torch.random.manual_seed(0)
|
|
# training_raw dir contains 300 folders from provided images and one folder with all images from these 300 folders,
|
|
# but augmented (sums up to about 60k images). (also see file image_augmentation)
|
|
train_data = RandomImagePixelationDataset(
|
|
r"training_raw",
|
|
width_range=(4, 32),
|
|
height_range=(4, 32),
|
|
size_range=(4, 16)
|
|
)
|
|
# eval_raw dir contains the remaining 50 folders for evaluation (~5k images)
|
|
eval_data = RandomImagePixelationDataset(
|
|
r"eval_raw",
|
|
width_range=(4, 32),
|
|
height_range=(4, 32),
|
|
size_range=(4, 16)
|
|
)
|
|
# working config: in_channels=1, out_channels=128, hidden_channels=6, kernel_size=7, dropout_rate=0.01
|
|
network = SimpleCNN(in_channels=1, out_channels=128, hidden_channels=6, kernel_size=7, dropout_rate=0.01)
|
|
|
|
model_param = filter(lambda p: p.requires_grad, network.parameters())
|
|
params = sum(np.prod(p.size()) for p in model_param)
|
|
print(f"Model has {params} parameters.")
|
|
|
|
epochs, train_loss, eval_loss = [], [], []
|
|
train_losses, eval_losses = training_loop(network, train_data, eval_data,
|
|
num_epochs=20, show_progress=True)
|
|
for epoch, (tl, el) in enumerate(zip(train_losses, eval_losses)):
|
|
print(f"Epoch: {epoch} --- Train loss: {tl:7.2f} --- Eval loss: {el:7.2f}")
|
|
epochs.append(epoch)
|
|
train_loss.append(tl)
|
|
eval_loss.append(el)
|
|
|
|
plt.plot(epochs, train_loss, label='train')
|
|
plt.plot(epochs, eval_loss, label='eval')
|
|
plt.xlabel('Epoch')
|
|
plt.ylabel('Loss')
|
|
plt.title('Loss per Epoch')
|
|
plt.show()
|