diff options
| -rw-r--r-- | .gitignore | 3 | ||||
| -rw-r--r-- | main.py | 48 | ||||
| -rw-r--r-- | trainer.py | 20 |
3 files changed, 60 insertions, 11 deletions
@@ -1,3 +1,6 @@ +# Tensorboard +runs/ + # Torch datas. data/ @@ -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}") @@ -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, |