From fc4617c9d8fd23c83fece9da914f9927ce5e5837 Mon Sep 17 00:00:00 2001 From: Steffen Illium Date: Thu, 18 Mar 2021 07:45:06 +0100 Subject: [PATCH] Final Train Runs --- metrics/multi_class_classification.py | 10 ++++++++-- modules/blocks.py | 28 ++++++++++++++++----------- modules/util.py | 6 +----- utils/callbacks.py | 23 ++++++++++++++++++++++ 4 files changed, 49 insertions(+), 18 deletions(-) create mode 100644 utils/callbacks.py diff --git a/metrics/multi_class_classification.py b/metrics/multi_class_classification.py index 0656376..6404821 100644 --- a/metrics/multi_class_classification.py +++ b/metrics/multi_class_classification.py @@ -26,10 +26,16 @@ class MultiClassScores(_BaseScores): ####################################################################################### # # INIT - y_true = torch.cat([output['batch_y'] for output in outputs]).cpu().numpy() + if isinstance(outputs['batch_y'], torch.Tensor): + y_true = outputs['batch_y'].cpu().numpy() + else: + y_true = torch.cat([output['batch_y'] for output in outputs]).cpu().numpy() y_true_one_hot = to_one_hot(y_true, self.model.params.n_classes) - y_pred = torch.cat([output['y'] for output in outputs]).squeeze().cpu().float().numpy() + if isinstance(outputs['y'], torch.Tensor): + y_pred = outputs['y'].cpu().numpy() + else: + y_pred = torch.cat([output['y'] for output in outputs]).squeeze().cpu().float().numpy() y_pred_max = np.argmax(y_pred, axis=1) class_names = {val: key for val, key in enumerate(class_names)} diff --git a/modules/blocks.py b/modules/blocks.py index 3bde036..ce3acf8 100644 --- a/modules/blocks.py +++ b/modules/blocks.py @@ -4,6 +4,7 @@ from pathlib import Path from typing import Union import torch +from performer_pytorch import FastAttention from torch import nn from torch.nn import functional as F @@ -72,7 +73,7 @@ class ConvModule(ShapeMixin, nn.Module): self.conv_kernel = conv_kernel # Modules - self.activation = activation() or F_x(None) + self.activation = activation() or nn.Identity() self.norm = nn.LayerNorm(self.in_shape, eps=1e-04) if use_norm else F_x(None) self.dropout = nn.Dropout2d(dropout) if dropout else F_x(None) self.pooling = nn.MaxPool2d(pooling_size) if pooling_size else F_x(None) @@ -232,16 +233,20 @@ class FeedForward(nn.Module): class Attention(nn.Module): - def __init__(self, dim, heads=8, dropout=0.): + def __init__(self, dim, heads=8, head_dim=64, dropout=0.): super().__init__() - self.heads = heads - self.scale = dim / heads ** -0.5 + inner_dim = head_dim * heads + project_out = not (heads == 1 and head_dim == dim) + + self.heads = heads + self.scale = head_dim ** -0.5 + + self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False) - self.to_qkv = nn.Linear(dim, dim * 3, bias=False) self.to_out = nn.Sequential( - nn.Linear(dim, dim), + nn.Linear(inner_dim, dim), nn.Dropout(dropout) - ) + ) if project_out else nn.Identity() def forward(self, x, mask=None, return_attn_weights=False): # noinspection PyTupleAssignmentBalance @@ -251,9 +256,9 @@ class Attention(nn.Module): q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), qkv) dots = torch.einsum('bhid,bhjd->bhij', q, k) * self.scale - mask_value = -torch.finfo(dots.dtype).max if mask is not None: + mask_value = -torch.finfo(dots.dtype).max mask = F.pad(mask.flatten(1), (1, 0), value=True) assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions' mask = mask[:, None, :] * mask[:, :, None] @@ -273,7 +278,7 @@ class Attention(nn.Module): class TransformerModule(ShapeMixin, nn.Module): - def __init__(self, in_shape, depth, heads, mlp_dim, dropout=None, use_norm=False, + def __init__(self, in_shape, depth, heads, mlp_dim, head_dim=32, dropout=None, use_norm=False, activation=nn.GELU, use_residual=True): super(TransformerModule, self).__init__() @@ -283,8 +288,9 @@ class TransformerModule(ShapeMixin, nn.Module): self.flat = Flatten(self.in_shape) if isinstance(self.in_shape, (tuple, list)) else F_x(in_shape) self.embedding_dim = self.flat.flat_shape - self.norm = nn.LayerNorm(self.embedding_dim) if use_norm else F_x(self.embedding_dim) - self.attns = nn.ModuleList([Attention(self.embedding_dim, heads=heads, dropout=dropout) for _ in range(depth)]) + self.norm = nn.LayerNorm(self.embedding_dim) if use_norm else F_x(None) + self.attns = nn.ModuleList([Attention(self.embedding_dim, heads=heads, dropout=dropout, head_dim=head_dim) + for _ in range(depth)]) self.mlps = nn.ModuleList([FeedForward(self.embedding_dim, mlp_dim, dropout=dropout, activation=activation) for _ in range(depth)]) diff --git a/modules/util.py b/modules/util.py index 2c29fb4..8b27fd0 100644 --- a/modules/util.py +++ b/modules/util.py @@ -207,15 +207,11 @@ class ShapeMixin: return shape -class F_x(ShapeMixin, nn.Module): +class F_x(ShapeMixin, nn.Identity): def __init__(self, in_shape): super(F_x, self).__init__() self.in_shape = in_shape - @staticmethod - def forward(x): - return x - class SlidingWindow(ShapeMixin, nn.Module): def __init__(self, in_shape, kernel, stride=1, padding=0, keepdim=False): diff --git a/utils/callbacks.py b/utils/callbacks.py new file mode 100644 index 0000000..70f410b --- /dev/null +++ b/utils/callbacks.py @@ -0,0 +1,23 @@ +from pytorch_lightning import Callback, Trainer, LightningModule + + +class BestScoresCallback(Callback): + + def __init__(self, *monitors) -> None: + super().__init__() + self.monitors = list(*monitors) + + self.best_scores = {monitor: 0.0 for monitor in self.monitors} + self.best_epoch = {monitor: 0 for monitor in self.monitors} + + def on_validation_end(self, trainer: Trainer, pl_module: LightningModule) -> None: + epoch = pl_module.current_epoch + + for monitor in self.best_scores.keys(): + current_score = trainer.callback_metrics.get(monitor) + if current_score is None: + pass + else: + self.best_scores[monitor] = max(self.best_scores[monitor], current_score) + if self.best_scores[monitor] == current_score: + self.best_epoch[monitor] = max(self.best_epoch[monitor], epoch)