modifed trainer to add callback and tensorboards

3 files changed, 60 insertions, 11 deletions

diff --git a/.gitignore b/.gitignore
index a0321c0..0c22bca 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,6 @@
+# Tensorboard
+runs/
+
 # Torch datas.
 data/
 
diff --git a/main.py b/main.py
index 7c7daa1..7473df1 100644
--- a/main.py
+++ b/main.py
@@ -53,15 +53,15 @@ class UNet(nn.Module):
         self.conv2 = nn.Conv2d(64, 64, 3, padding=1)
         self.maxpool1 = nn.MaxPool2d(2, 2)
         self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
-        self.att1 = SelfAttention(128, 8)
+        # self.att1 = SelfAttention(128, 8)
         self.conv4 = nn.Conv2d(128, 128, 3, padding=1)
         self.maxpool2 = nn.MaxPool2d(2, 2)
         self.conv5 = nn.Conv2d(128, 256, 3, padding=1)
-        self.att2 = SelfAttention(256, 8)
+        # self.att2 = SelfAttention(256, 8)
         self.conv6 = nn.Conv2d(256, 256, 3, padding=1)
         self.upconv1 = nn.ConvTranspose2d(256, 128, 2, stride=2)
         self.conv7 = nn.Conv2d(256, 128, 3, padding=1)
-        self.att3 = SelfAttention(128, 8)
+        # self.att3 = SelfAttention(128, 8)
         self.conv8 = nn.Conv2d(128, 128, 3, padding=1)
         self.upconv2 = nn.ConvTranspose2d(128, 64, 2, stride=2)
         self.conv9 = nn.Conv2d(128, 64, 3, padding=1)
@@ -83,16 +83,16 @@ class UNet(nn.Module):
         x1 = F.relu(self.conv2(x1))
         x2 = self.maxpool1(x1)
         x2 = F.relu(self.conv3(x2))
-        x2 = self.att1(x2)
+        # x2 = self.att1(x2)
         x2 = F.relu(self.conv4(x2))
         x3 = self.maxpool2(x2)
         x3 = F.relu(self.conv5(x3))
-        x3 = self.att2(x3)
+        # x3 = self.att2(x3)
         x5 = F.relu(self.conv6(x3))
         x6 = self.upconv1(x5)
         x6 = torch.cat((x6, x2), dim=1)
         x6 = F.relu(self.conv7(x6))
-        x6 = self.att3(x6)
+        # x6 = self.att3(x6)
         x6 = F.relu(self.conv8(x6))
         x7 = self.upconv2(x6)
         x7 = torch.cat((x7, x1), dim=1)
@@ -262,10 +262,39 @@ if autoencoder is not None:
 NB_CHANNEL, IMG_SIZE, _ = img.shape
 NB_LABEL = 1
 
-EPOCHS = 0
+EPOCHS = 200
 LEARNING_RATE = 2e-4
 
 
+def epoch_callback(epoch_i, epochs, model, trainer):
+    if epoch_i % 10 == 0 or epoch_i == epochs - 1:
+        print("Calculating metrics...")
+        with torch.no_grad():
+            batch_size = 64
+            n_batches = 16
+            n_samples = batch_size * n_batches
+
+            fakes = np.zeros((0, NB_CHANNEL, IMG_SIZE, IMG_SIZE))
+            for _ in range(n_batches):
+                x = torch.randn(batch_size, NB_CHANNEL, IMG_SIZE, IMG_SIZE).to(DEVICE)
+                vec = torch.randint(0, NB_LABEL, (batch_size,)).to(DEVICE)
+                vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(device=DEVICE, dtype=torch.float32)
+
+                for t in range(DIFFU_STEPS, 0, -1):
+                    t_tensor = torch.tensor([[t]] * batch_size, device=DEVICE, dtype=torch.float32)
+                    x = p_xt_1_xt(model, x, t_tensor, vec)
+
+                fakes = np.concatenate((fakes, x.cpu().numpy()))
+
+            reals = torch.stack([dataset[i][0] for i in range(n_samples)]).cpu().numpy()
+            reals = reals * 2 - 1
+
+            trainer.writer.add_scalars('Metrics/FID', {trainer.date_time: fid(reals, fakes)}, epoch_i)
+            trainer.writer.add_scalars('Metrics/KL', {trainer.date_time: kl(reals, fakes)}, epoch_i)
+            trainer.writer.add_scalars('Metrics/RKL', {trainer.date_time: kl(fakes, reals)}, epoch_i)
+            trainer.writer.add_scalars('Metrics/JSD', {trainer.date_time: jsd(reals, fakes)}, epoch_i)
+
+
 if __name__ == '__main__':
 
     ############
@@ -277,7 +306,7 @@ if __name__ == '__main__':
     # Load the model.
     model = UNet().to(DEVICE)
     try:
-        model.load_state_dict(torch.load('edf_att_model.pth'))
+        model.load_state_dict(torch.load('model.pth'))
     except FileNotFoundError:
         print("No model found, training a new one.")
         pass
@@ -294,7 +323,7 @@ if __name__ == '__main__':
     epochs = EPOCHS
 
     # Train the model.
-    trainer.train(model, train_loader, epochs, optimizer, criterion)
+    trainer.train(model, train_loader, epochs, optimizer, criterion, epoch_callbacks=[epoch_callback])
 
     # Save the model.
     torch.save(model.state_dict(), 'model.pth')
@@ -419,7 +448,6 @@ if __name__ == '__main__':
         fid_score = fid(reals, fakes)
         print(f"FID score: {fid_score}")
 
-
         kl_score = kl(reals, fakes)
         print(f"KL divergence: {kl_score}")
 
diff --git a/trainer.py b/trainer.py
index 7854b63..f088e35 100644
--- a/trainer.py
+++ b/trainer.py
@@ -1,10 +1,13 @@
 import torch
 import torch.utils.data
+from torch.utils.tensorboard import SummaryWriter
 
 import rich.progress
 
 from typing import *
 
+import datetime
+
 
 class TrainProgress(rich.progress.Progress):
     """A progress bar which tracks the progress of training epochs."""
@@ -289,6 +292,7 @@ class Trainer:
     def __init__(self):
         """Initialize the trainer."""
         self.progress: TrainProgress | None = None
+        self.writer: SummaryWriter | None = None
 
     def train(
         self: 'Trainer',
@@ -301,6 +305,7 @@ class Trainer:
         test_loader: torch.utils.data.DataLoader | None = None,
         metrics: List[Callable[[torch.Tensor,
                                 torch.Tensor], torch.Tensor]] = [],
+        epoch_callbacks: List[Callable[[int, torch.nn.Module], None]] = [],
     ):
         """Train the model for the given number of epochs.
 
@@ -311,7 +316,12 @@ class Trainer:
             optimizer (torch.optim.Optimizer): The optimizer to use.
             criterion (Callable[[torch.Tensor, torch.Tensor], torch.Tensor]): The loss function to use.
             val_loader (torch.utils.data.DataLoader, optional): The validation dataset. Defaults to None.
+            test_loader (torch.utils.data.DataLoader, optional): The test dataset. Defaults to None.
+            metrics (List[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]], optional): The metrics to use. Defaults to [].
+            epoch_callbacks (List[Callable[[int, torch.nn.Module], None]], optional): The callbacks to call at the end of each epoch. Defaults to [].
         """
+        self.writer = SummaryWriter(log_dir='runs')
+        self.date_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
         with TrainProgress(
             nb_epochs=epochs,
             train_size=len(train_loader),
@@ -320,13 +330,14 @@ class Trainer:
         ) as progress:
             self.progress = progress
 
-            for _ in range(epochs):
+            for epoch_i in range(epochs):
                 self.train_epoch(
                     model,
                     train_loader,
                     optimizer,
                     criterion,
                     metrics,
+                    epoch_i,
                 )
                 if val_loader:
                     self.validate(
@@ -334,12 +345,15 @@ class Trainer:
                         val_loader,
                         metrics + [criterion],
                     )
+                for callback in epoch_callbacks:
+                    callback(epoch_i=epoch_i, epochs=epochs, model=model, trainer=self)
             if test_loader:
                 self.test(
                     model,
                     test_loader,
                     metrics + [criterion],
                 )
+        self.writer.close()
 
     def train_epoch(
         self: 'Trainer',
@@ -348,6 +362,7 @@ class Trainer:
         optimizer: torch.optim.Optimizer,
         criterion: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
         metrics: list[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]],
+        epoch_i: int,
     ):
         """Train the model for one epoch.
 
@@ -357,6 +372,7 @@ class Trainer:
             optimizer (torch.optim.Optimizer): The optimizer to use.
             criterion (Callable[[torch.Tensor, torch.Tensor], torch.Tensor]): The loss function to use.
             metrics (list[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]]): The metrics to use.
+            epoch_i (int): The current epoch.
         """
         model.train()
         for batch in train_loader:
@@ -378,6 +394,8 @@ class Trainer:
             self.progress.step()
             self.progress.new_train_values(values)
 
+        self.writer.add_scalars('Criterion/train', {self.date_time: loss.item()}, epoch_i)
+
     def validate(
         self: 'Trainer',
         model: torch.nn.Module,