Shortcuts

Source code for catalyst.contrib.nn.modules.lama

import torch
from torch import nn

from catalyst.utils import outer_init


[docs]class TemporalLastPooling(nn.Module):
[docs] def forward(self, x: torch.Tensor, mask: torch.Tensor = None): x_out = x[:, -1:, :] return x_out
[docs]class TemporalAvgPooling(nn.Module):
[docs] def forward(self, x: torch.Tensor, mask: torch.Tensor = None): if mask is None: x_out = x.mean(1, keepdim=True) else: x_ = torch.sum(x * mask.float(), dim=1, keepdim=True) mask_ = torch.sum(mask.float(), dim=1, keepdim=True) x_out = x_ / mask_ return x_out
[docs]class TemporalMaxPooling(nn.Module):
[docs] def forward(self, x: torch.Tensor, mask: torch.Tensor = None): if mask is not None: mask_ = (~mask.bool()).float() * (-x.max()).float() x = torch.sum(x + mask_, dim=1, keepdim=True) x_out = x.max(1, keepdim=True)[0] return x_out
[docs]class TemporalAttentionPooling(nn.Module): name2activation = { "softmax": nn.Softmax(dim=1), "tanh": nn.Tanh(), "sigmoid": nn.Sigmoid() } def __init__(self, in_features, activation=None, kernel_size=1, **params): super().__init__() self.in_features = in_features activation = activation or "softmax" self.attention_pooling = nn.Sequential( nn.Conv1d( in_channels=in_features, out_channels=1, kernel_size=kernel_size, **params ), TemporalAttentionPooling.name2activation[activation] ) self.attention_pooling.apply(outer_init)
[docs] def forward(self, x: torch.Tensor, mask: torch.Tensor = None): """ :param x: [batch_size, history_len, feature_size] :return: """ batch_size, history_len, feature_size = x.shape x = x.view(batch_size, history_len, -1) x_a = x.transpose(1, 2) x_attn = (self.attention_pooling(x_a) * x_a).transpose(1, 2) x_attn = x_attn.sum(1, keepdim=True) return x_attn
[docs]class TemporalConcatPooling(nn.Module): def __init__(self, in_features, history_len=1): super().__init__() self.in_features = in_features self.out_features = in_features * history_len
[docs] def forward(self, x: torch.Tensor, mask: torch.Tensor = None): """ :param x: [batch_size, history_len, feature_size] :return: """ x = x.view(x.shape[0], -1) return x
[docs]class TemporalDropLastWrapper(nn.Module): def __init__(self, net): super().__init__() self.net = net
[docs] def forward(self, x: torch.Tensor, mask: torch.Tensor = None): x = x[:, :-1, :] x_out = self.net(x) return x_out
def get_pooling(key, in_features, **params): key_ = key.split("_", 1)[0] if key_ == "last": return TemporalLastPooling() elif key_ == "avg": layer = TemporalAvgPooling() elif key_ == "max": layer = TemporalMaxPooling() elif key_ in ["softmax", "tanh", "sigmoid"]: layer = TemporalAttentionPooling( in_features=in_features, activation=key_, **params ) else: raise NotImplementedError() if "droplast" in key: layer = TemporalDropLastWrapper(layer) return layer
[docs]class LamaPooling(nn.Module): available_groups = [ "last", "avg", "avg_droplast", "max", "max_droplast", "sigmoid", "sigmoid_droplast", "softmax", "softmax_droplast", "tanh", "tanh_droplast", ] def __init__(self, in_features, groups=None): super().__init__() self.in_features = in_features self.groups = groups \ or ["last", "avg_droplast", "max_droplast", "softmax_droplast"] self.out_features = in_features * len(self.groups) groups = {} for key in self.groups: if isinstance(key, str): groups[key] = get_pooling(key, self.in_features) elif isinstance(key, dict): key_ = key.pop("key") groups[key_] = get_pooling(key_, in_features, **key) else: raise NotImplementedError() self.groups = nn.ModuleDict(groups)
[docs] def forward(self, x: torch.Tensor, mask: torch.Tensor = None): """ :param x: [batch_size, history_len, feature_size] :return: """ batch_size, history_len, feature_size = x.shape x_ = [] for pooling_fn in self.groups.values(): features_ = pooling_fn(x, mask) x_.append(features_) x = torch.cat(x_, dim=1) x = x.view(batch_size, -1) return x
__all__ = [ "TemporalLastPooling", "TemporalAvgPooling", "TemporalMaxPooling", "TemporalDropLastWrapper", "TemporalAttentionPooling", "TemporalConcatPooling", "LamaPooling", ]