Transformer running

This commit is contained in:
Steffen Illium
2021-03-04 12:01:08 +01:00
parent b5e3e5aec1
commit f89f0f8528
14 changed files with 349 additions and 80 deletions

View File

@@ -22,8 +22,8 @@ DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
###################
class LinearModule(ShapeMixin, nn.Module):
def __init__(self, in_shape, out_features, bias=True, activation=None,
norm=False, dropout: Union[int, float] = 0, **kwargs):
def __init__(self, in_shape, out_features, use_bias=True, activation=None,
use_norm=False, dropout: Union[int, float] = 0, **kwargs):
if list(kwargs.keys()):
warnings.warn(f'The following arguments have been ignored: \n {list(kwargs.keys())}')
super(LinearModule, self).__init__()
@@ -31,8 +31,8 @@ class LinearModule(ShapeMixin, nn.Module):
self.in_shape = in_shape
self.flat = Flatten(self.in_shape) if isinstance(self.in_shape, (tuple, list)) else F_x(in_shape)
self.dropout = nn.Dropout(dropout) if dropout else F_x(self.flat.shape)
self.norm = nn.BatchNorm1d(self.flat.shape) if norm else F_x(self.flat.shape)
self.linear = nn.Linear(self.flat.shape, out_features, bias=bias)
self.norm = nn.LayerNorm(self.flat.shape) if use_norm else F_x(self.flat.shape)
self.linear = nn.Linear(self.flat.shape, out_features, bias=use_bias)
self.activation = activation() if activation else F_x(self.linear.out_features)
def forward(self, x):
@@ -47,13 +47,14 @@ class LinearModule(ShapeMixin, nn.Module):
class ConvModule(ShapeMixin, nn.Module):
def __init__(self, in_shape, conv_filters, conv_kernel, activation: nn.Module = nn.ELU, pooling_size=None,
bias=True, norm=False, dropout: Union[int, float] = 0, trainable: bool = True,
bias=True, use_norm=False, dropout: Union[int, float] = 0, trainable: bool = True,
conv_class=nn.Conv2d, conv_stride=1, conv_padding=0, **kwargs):
super(ConvModule, self).__init__()
assert isinstance(in_shape, (tuple, list)), f'"in_shape" should be a [list, tuple], but was {type(in_shape)}'
assert len(in_shape) == 3, f'Length should be 3, but was {len(in_shape)}'
warnings.warn(f'The following arguments have been ignored: \n {list(kwargs.keys())}')
if norm and not trainable:
if len(kwargs.keys()):
warnings.warn(f'The following arguments have been ignored: \n {list(kwargs.keys())}')
if use_norm and not trainable:
warnings.warn('You set this module to be not trainable but the running norm is active.\n' +
'We set it to "eval" mode.\n' +
'Keep this in mind if you do a finetunning or retraining step.'
@@ -72,9 +73,9 @@ class ConvModule(ShapeMixin, nn.Module):
# Modules
self.activation = activation() or F_x(None)
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)
self.norm = nn.BatchNorm2d(in_channels, eps=1e-04) if norm else F_x(None)
self.conv = conv_class(in_channels, self.conv_filters, self.conv_kernel, bias=bias,
padding=self.padding, stride=self.stride
)
@@ -134,7 +135,7 @@ class DeConvModule(ShapeMixin, nn.Module):
def __init__(self, in_shape, conv_filters, conv_kernel, conv_stride=1, conv_padding=0,
dropout: Union[int, float] = 0, autopad=0,
activation: Union[None, nn.Module] = nn.ReLU, interpolation_scale=0,
bias=True, norm=False, **kwargs):
bias=True, use_norm=False, **kwargs):
super(DeConvModule, self).__init__()
warnings.warn(f'The following arguments have been ignored: \n {list(kwargs.keys())}')
in_channels, height, width = in_shape[0], in_shape[1], in_shape[2]
@@ -146,7 +147,7 @@ class DeConvModule(ShapeMixin, nn.Module):
self.autopad = AutoPad() if autopad else lambda x: x
self.interpolation = Interpolate(scale_factor=interpolation_scale) if interpolation_scale else lambda x: x
self.norm = nn.BatchNorm2d(in_channels, eps=1e-04) if norm else F_x(self.in_shape)
self.norm = nn.LayerNorm(in_channels, eps=1e-04) if use_norm else F_x(self.in_shape)
self.dropout = nn.Dropout2d(dropout) if dropout else F_x(self.in_shape)
self.de_conv = nn.ConvTranspose2d(in_channels, self.conv_filters, self.conv_kernel, bias=bias,
padding=self.padding, stride=self.stride)
@@ -166,14 +167,13 @@ class DeConvModule(ShapeMixin, nn.Module):
class ResidualModule(ShapeMixin, nn.Module):
def __init__(self, in_shape, module_class, n, norm=False, **module_parameters):
def __init__(self, in_shape, module_class, n, use_norm=False, **module_parameters):
assert n >= 1
super(ResidualModule, self).__init__()
self.in_shape = in_shape
module_parameters.update(in_shape=in_shape)
if norm:
norm = nn.BatchNorm1d if len(self.in_shape) <= 2 else nn.BatchNorm2d
self.norm = norm(self.in_shape if isinstance(self.in_shape, int) else self.in_shape[0])
if use_norm:
self.norm = nn.LayerNorm(self.in_shape if isinstance(self.in_shape, int) else self.in_shape[0])
else:
self.norm = F_x(self.in_shape)
self.activation = module_parameters.get('activation', None)
@@ -216,13 +216,14 @@ class RecurrentModule(ShapeMixin, nn.Module):
class FeedForward(nn.Module):
def __init__(self, dim, hidden_dim, dropout=0.):
def __init__(self, dim, hidden_dim, dropout=0., activation=nn.GELU):
super().__init__()
self.net = nn.Sequential(
nn.Linear(dim, hidden_dim),
nn.GELU(),
activation() or F_x(None),
nn.Dropout(dropout),
nn.Linear(hidden_dim, dim),
activation() or F_x(None),
nn.Dropout(dropout)
)
@@ -272,18 +273,20 @@ class Attention(nn.Module):
class TransformerModule(ShapeMixin, nn.Module):
def __init__(self, in_shape, depth, heads, mlp_dim, dropout=None, use_norm=False, activation='gelu'):
def __init__(self, in_shape, depth, heads, mlp_dim, dropout=None, use_norm=False,
activation=nn.GELU, use_residual=True):
super(TransformerModule, self).__init__()
self.in_shape = in_shape
self.use_residual = use_residual
self.flat = Flatten(self.in_shape) if isinstance(self.in_shape, (tuple, list)) else F_x(in_shape)
self.layers = nn.ModuleList([])
self.embedding_dim = self.flat.flat_shape
self.norm = nn.LayerNorm(self.embedding_dim)
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.mlps = nn.ModuleList([FeedForward(self.embedding_dim, mlp_dim, dropout=dropout) for _ in range(depth)])
self.mlps = nn.ModuleList([FeedForward(self.embedding_dim, mlp_dim, dropout=dropout, activation=activation)
for _ in range(depth)])
def forward(self, x, mask=None, return_attn_weights=False, **_):
tensor = self.flat(x)
@@ -297,11 +300,11 @@ class TransformerModule(ShapeMixin, nn.Module):
attn_weights.append(attn_weight)
else:
attn_tensor = attn(attn_tensor, mask=mask)
tensor = attn_tensor + tensor
tensor = tensor + attn_tensor if self.use_residual else attn_tensor
# MLP
mlp_tensor = self.norm(tensor)
mlp_tensor = mlp(mlp_tensor)
tensor = tensor + mlp_tensor
tensor = tensor + mlp_tensor if self.use_residual else mlp_tensor
return (tensor, attn_weights) if return_attn_weights else tensor

View File

@@ -183,10 +183,11 @@ class BaseCNNEncoder(ShapeMixin, nn.Module):
# noinspection PyUnresolvedReferences
def __init__(self, in_shape, lat_dim=256, use_bias=True, use_norm=False, dropout: Union[int, float] = 0,
latent_activation: Union[nn.Module, None] = None, activation: nn.Module = nn.ELU,
filters: List[int] = None, kernels: List[int] = None, **kwargs):
filters: List[int] = None, kernels: Union[List[int], int, None] = None, **kwargs):
super(BaseCNNEncoder, self).__init__()
assert filters, '"Filters" has to be a list of int'
assert kernels, '"Kernels" has to be a list of int'
kernels = kernels or [3] * len(filters)
kernels = kernels if not isinstance(kernels, int) else [kernels] * len(filters)
assert len(kernels) == len(filters), 'Length of "Filters" and "Kernels" has to be same.'
# Optional Padding for odd image-sizes

View File

@@ -1,7 +1,5 @@
import inspect
from argparse import ArgumentParser
from functools import reduce
from matplotlib import pyplot as plt
from abc import ABC
from pathlib import Path
@@ -12,14 +10,77 @@ from pytorch_lightning.utilities import argparse_utils
from torch import nn
from torch.nn import functional as F, Unfold
from sklearn.metrics import ConfusionMatrixDisplay
# Utility - Modules
###################
from ..utils.model_io import ModelParameters
from ..utils.tools import locate_and_import_class, add_argparse_args
from ..utils.tools import add_argparse_args
try:
import pytorch_lightning as pl
class PLMetrics(pl.metrics.Metric):
def __init__(self, n_classes, tag=''):
super(PLMetrics, self).__init__()
self.n_classes = n_classes
self.tag = tag
self.accuracy_score = pl.metrics.Accuracy(compute_on_step=False)
self.precision = pl.metrics.Precision(num_classes=self.n_classes, average='macro', compute_on_step=False)
self.recall = pl.metrics.Recall(num_classes=self.n_classes, average='macro', compute_on_step=False)
self.confusion_matrix = pl.metrics.ConfusionMatrix(self.n_classes, normalize='true', compute_on_step=False)
# self.precision_recall_curve = pl.metrics.PrecisionRecallCurve(self.n_classes, compute_on_step=False)
# self.average_prec = pl.metrics.AveragePrecision(self.n_classes, compute_on_step=True)
# self.roc = pl.metrics.ROC(self.n_classes, compute_on_step=False)
self.fbeta = pl.metrics.FBeta(self.n_classes, average='macro', compute_on_step=False)
self.f1 = pl.metrics.F1(self.n_classes, average='macro', compute_on_step=False)
def __iter__(self):
return iter(((name, metric) for name, metric in self._modules.items()))
def update(self, preds, target) -> None:
for _, metric in self:
metric.update(preds, target)
def reset(self) -> None:
for _, metric in self:
metric.reset()
def compute(self) -> dict:
tag = f'{self.tag}_' if self.tag else ''
return {f'{tag}{metric_name}_score': metric.compute() for metric_name, metric in self}
def compute_and_prepare(self):
pl_metrics = self.compute()
images_from_metrics = dict()
for metric_name in list(pl_metrics.keys()):
if 'curve' in metric_name:
continue
roc_curve = pl_metrics.pop(metric_name)
print('debug_point')
elif 'matrix' in metric_name:
matrix = pl_metrics.pop(metric_name)
fig1, ax1 = plt.subplots(dpi=96)
disp = ConfusionMatrixDisplay(confusion_matrix=matrix.cpu().numpy(),
display_labels=[i for i in range(self.n_classes)]
)
disp.plot(include_values=True, ax=ax1)
images_from_metrics[metric_name] = fig1
elif 'ROC' in metric_name:
continue
roc = pl_metrics.pop(metric_name)
print('debug_point')
else:
pl_metrics[metric_name] = pl_metrics[metric_name].cpu().item()
return pl_metrics, images_from_metrics
class LightningBaseModule(pl.LightningModule, ABC):
@classmethod
@@ -49,6 +110,9 @@ try:
self._weight_init = weight_init
self.params = ModelParameters(model_parameters)
self.metrics = PLMetrics(self.params.n_classes, tag='PL')
pass
def size(self):
return self.shape