upload of files

2023-12-09 17:28:07 +01:00
parent f6daee7e40
commit 454acefe1e
14 changed files with 597 additions and 0 deletions
--- a/training_loop.py
+++ b/training_loop.py
@@ -0,0 +1,131 @@
+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()