opensportslib 0.0.1.dev2__py3-none-any.whl

opensportslib/models/heads/builder.py

@@ -0,0 +1,266 @@
+"""
+Copyright 2022 James Hong, Haotian Zhang, Matthew Fisher, Michael Gharbi,
+Kayvon Fatahalian
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+this list of conditions and the following disclaimer in the documentation and/or
+other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors
+may be used to endorse or promote products derived from this software without
+specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
+ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+"""
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+
+from opensportslib.models.utils.modules import *
+
+
+def build_head(cfg, default_args=None):
+    """Build a head from config dict.
+
+    Args:
+        cfg (dict): Config dict. It should at least contain the key "type".
+        default_args (dict | None, optional): Default initialization arguments.
+            Default: None.
+
+    Returns:
+        head: The constructed head.
+    """
+    if cfg.type == "TrackingClassifier":
+        head = TrackingClassifierHead(
+            input_dim=default_args["input_dim"],
+            hidden_dim=cfg.hidden_dim,
+            num_classes=cfg.num_classes,
+            dropout=cfg.dropout
+        )
+    elif cfg.type == "LinearLayer":
+        head = LinearLayerHead(input_dim=cfg.input_dim, output_dim=cfg.num_classes + 1)
+    elif cfg.type == "SpottingCALF":
+        head = SpottingCALFHead(
+            num_classes=cfg.num_classes,
+            dim_capsule=cfg.dim_capsule,
+            num_detections=cfg.num_detections,
+            chunk_size=cfg.chunk_size,
+        )
+    elif cfg.type in ["", "gru", "deeper_gru", "mstcn", "asformer"]:
+        head = TemporalE2EHead(cfg.type, cfg.feat_dim, cfg.num_classes)
+    elif cfg.type == "MV_LinearLayer":
+        head = MVHead(input_dim=cfg.feat_dim, output_dim=cfg.num_classes)
+    else:
+        head = None
+
+    return head
+
+
+class MVHead(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+        print("Inside MVHead")
+        self.fc = nn.Sequential(
+            nn.LayerNorm(input_dim),
+            nn.Linear(input_dim, input_dim),
+            nn.Linear(input_dim, output_dim)
+        )
+
+    def forward(self, x):
+        pred = self.fc(x)
+        return pred
+
+class TrackingClassifierHead(nn.Module):
+    """Classification head for tracking data."""
+    
+    def __init__(self, input_dim, hidden_dim, num_classes, dropout=0.1):
+        super().__init__()
+        
+        self.classifier = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(input_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout * 0.5),
+            nn.Linear(hidden_dim, num_classes)
+        )
+    
+    def forward(self, x):
+        return self.classifier(x)
+    
+class TemporalE2EHead(nn.Module):
+    def __init__(self, temporal_arch, feat_dim, num_classes):
+        super().__init__()
+        # Prevent the GRU params from going too big (cap it at a RegNet-Y 800MF)
+        MAX_GRU_HIDDEN_DIM = 768
+        if "gru" in temporal_arch:
+            hidden_dim = feat_dim
+            if hidden_dim > MAX_GRU_HIDDEN_DIM:
+                hidden_dim = MAX_GRU_HIDDEN_DIM
+                print("Clamped GRU hidden dim: {} -> {}".format(feat_dim, hidden_dim))
+            if temporal_arch in ("gru", "deeper_gru"):
+                self._pred_fine = GRUPrediction(
+                    feat_dim,
+                    num_classes,
+                    hidden_dim,
+                    num_layers=3 if temporal_arch[0] == "d" else 1,
+                )
+            else:
+                raise NotImplementedError(temporal_arch)
+        elif temporal_arch == "mstcn":
+            self._pred_fine = TCNPrediction(feat_dim, num_classes, 3)
+        elif temporal_arch == "asformer":
+            self._pred_fine = ASFormerPrediction(feat_dim, num_classes, 3)
+        elif temporal_arch == "":
+            self._pred_fine = FCPrediction(feat_dim, num_classes)
+
+    def forward(self, inputs):
+        return self._pred_fine(inputs)
+
+
+class LinearLayerHead(torch.nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super(LinearLayerHead, self).__init__()
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+
+        self.drop = torch.nn.Dropout(p=0.4)
+        self.head = torch.nn.Linear(input_dim, output_dim)
+        self.sigm = torch.nn.Sigmoid()
+
+    def forward(self, inputs):
+        return self.sigm(self.head(self.drop(inputs)))
+
+
+class SpottingCALFHead(torch.nn.Module):
+    def __init__(self, num_classes, dim_capsule, num_detections, chunk_size):
+        super(SpottingCALFHead, self).__init__()
+
+        self.num_classes = num_classes
+        self.dim_capsule = dim_capsule
+        self.num_detections = num_detections
+        self.chunk_size = chunk_size
+
+        # -------------------
+        # detection module
+        # -------------------
+        self.max_pool_spot = nn.MaxPool2d(kernel_size=(3, 1), stride=(2, 1))
+        self.kernel_spot_size = 3
+        self.pad_spot_1 = nn.ZeroPad2d(
+            (
+                0,
+                0,
+                (self.kernel_spot_size - 1) // 2,
+                self.kernel_spot_size - 1 - (self.kernel_spot_size - 1) // 2,
+            )
+        )
+        self.conv_spot_1 = nn.Conv2d(
+            in_channels=num_classes * (dim_capsule + 1),
+            out_channels=32,
+            kernel_size=(self.kernel_spot_size, 1),
+        )
+        self.max_pool_spot_1 = nn.MaxPool2d(kernel_size=(3, 1), stride=(2, 1))
+        self.pad_spot_2 = nn.ZeroPad2d(
+            (
+                0,
+                0,
+                (self.kernel_spot_size - 1) // 2,
+                self.kernel_spot_size - 1 - (self.kernel_spot_size - 1) // 2,
+            )
+        )
+        self.conv_spot_2 = nn.Conv2d(
+            in_channels=32, out_channels=16, kernel_size=(self.kernel_spot_size, 1)
+        )
+        self.max_pool_spot_2 = nn.MaxPool2d(kernel_size=(3, 1), stride=(2, 1))
+
+        # Confidence branch
+        self.conv_conf = nn.Conv2d(
+            in_channels=16 * (chunk_size // 8 - 1),
+            out_channels=self.num_detections * 2,
+            kernel_size=(1, 1),
+        )
+
+        # Class branch
+        self.conv_class = nn.Conv2d(
+            in_channels=16 * (chunk_size // 8 - 1),
+            out_channels=self.num_detections * self.num_classes,
+            kernel_size=(1, 1),
+        )
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, conv_seg, output_segmentation):
+        # ---------------
+        # Spotting module
+        # ---------------
+
+        # Concatenation of the segmentation score to the capsules
+        output_segmentation_reverse = 1 - output_segmentation
+        # print("Output_segmentation_reverse size: ", output_segmentation_reverse.size())
+
+        output_segmentation_reverse_reshaped = output_segmentation_reverse.unsqueeze(2)
+        # print("Output_segmentation_reverse_reshaped size: ", output_segmentation_reverse_reshaped.size())
+
+        output_segmentation_reverse_reshaped_permutted = (
+            output_segmentation_reverse_reshaped.permute(0, 3, 1, 2)
+        )
+        # print("Output_segmentation_reverse_reshaped_permutted size: ", output_segmentation_reverse_reshaped_permutted.size())
+
+        concatenation_2 = torch.cat(
+            (conv_seg, output_segmentation_reverse_reshaped_permutted), dim=1
+        )
+        # print("Concatenation_2 size: ", concatenation_2.size())
+
+        conv_spot = self.max_pool_spot(F.relu(concatenation_2))
+        # print("Conv_spot size: ", conv_spot.size())
+
+        conv_spot_1 = F.relu(self.conv_spot_1(self.pad_spot_1(conv_spot)))
+        # print("Conv_spot_1 size: ", conv_spot_1.size())
+
+        conv_spot_1_pooled = self.max_pool_spot_1(conv_spot_1)
+        # print("Conv_spot_1_pooled size: ", conv_spot_1_pooled.size())
+
+        conv_spot_2 = F.relu(self.conv_spot_2(self.pad_spot_2(conv_spot_1_pooled)))
+        # print("Conv_spot_2 size: ", conv_spot_2.size())
+
+        conv_spot_2_pooled = self.max_pool_spot_2(conv_spot_2)
+        # print("Conv_spot_2_pooled size: ", conv_spot_2_pooled.size())
+
+        spotting_reshaped = conv_spot_2_pooled.view(
+            conv_spot_2_pooled.size()[0], -1, 1, 1
+        )
+        # print("Spotting_reshape size: ", spotting_reshaped.size())
+
+        # Confindence branch
+        conf_pred = torch.sigmoid(
+            self.conv_conf(spotting_reshaped).view(
+                spotting_reshaped.shape[0], self.num_detections, 2
+            )
+        )
+        # print("Conf_pred size: ", conf_pred.size())
+
+        # Class branch
+        conf_class = self.softmax(
+            self.conv_class(spotting_reshaped).view(
+                spotting_reshaped.shape[0], self.num_detections, self.num_classes
+            )
+        )
+        # print("Conf_class size: ", conf_class.size())
+
+        output_spotting = torch.cat((conf_pred, conf_class), dim=-1)
+        # print("Output_spotting size: ", output_spotting.size())
+
+        return output_spotting

opensportslib/models/neck/builder.py

@@ -0,0 +1,210 @@
+import torch
+from torch import nn
+from opensportslib.core.utils.data import batch_tensor, unbatch_tensor
+
+def build_neck(cfg, default_args=None):
+    if cfg.type == "MV_Aggregate":
+        neck = MVAggregate(
+            agr_type=cfg.agr_type,
+            model=default_args["model"],
+            feat_dim=default_args["feat_dim"],
+            lifting_net=default_args["lifting_net"] if "lifting_net" in default_args else nn.Sequential()
+        )
+    
+    elif cfg.type == "TemporalAggregation":
+        neck = TemporalAggregation(
+            temporal_type=cfg.agr_type,
+            hidden_dim=cfg.hidden_dim,
+            window_size=default_args["window_size"],
+            dropout=cfg.dropout,
+            use_position_encoding=getattr(cfg, "use_position_encoding", False),
+            num_attention_heads=getattr(cfg, "num_attention_heads", 4),
+            lstm_dropout=getattr(cfg, "lstm_dropout", 0.1)
+        )
+    else:
+        raise ValueError(f"Unknown neck type: {cfg.type}")
+    return neck
+
+
+class MVAggregate(nn.Module):
+    def __init__(self, agr_type, model, feat_dim, lifting_net=nn.Sequential()):
+        super().__init__()
+        self.agr_type = agr_type
+        self.model = model
+        self.feat_dim = feat_dim
+        self.lifting_net = lifting_net
+        print("Inside NECK BUILDER - AGR TYPE:", self.agr_type)
+
+        if self.agr_type == "max":
+            self.aggregation_model = ViewMaxAggregate(model=model, lifting_net=lifting_net)
+        elif self.agr_type == "mean":
+            self.aggregation_model = ViewAvgAggregate(model=model, lifting_net=lifting_net)
+        else:
+            # avg
+            self.aggregation_model = WeightedAggregate(model=model, feat_dim=feat_dim, lifting_net=lifting_net)
+
+        self.inter = nn.Sequential(
+            nn.LayerNorm(feat_dim),
+            nn.Linear(feat_dim, feat_dim),
+            nn.Linear(feat_dim, feat_dim),
+        )
+    
+    def forward(self, mvimages):
+        pooled_view, attention = self.aggregation_model(mvimages)
+        inter = self.inter(pooled_view)
+        return inter, attention
+
+        
+
+class WeightedAggregate(nn.Module):
+    def __init__(self,  model, feat_dim, lifting_net=nn.Sequential()):
+        super().__init__()
+        self.model = model
+        self.lifting_net = lifting_net
+        self.feature_dim = feat_dim
+
+        r1 = -1
+        r2 = 1
+        self.attention_weights = nn.Parameter((r1 - r2) * torch.rand(feat_dim, feat_dim) + r2)
+
+        self.normReLu = nn.Sequential(
+            nn.LayerNorm(feat_dim),
+            nn.ReLU()
+        )        
+
+        self.relu = nn.ReLU()
+   
+
+
+    def forward(self, mvimages):
+        B, V, C, D, H, W = mvimages.shape # Batch, Views, Channel, Depth, Height, Width
+        aux = self.lifting_net(unbatch_tensor(self.model(batch_tensor(mvimages, dim=1, squeeze=True)), B, dim=1, unsqueeze=True))
+
+
+        ##################### VIEW ATTENTION #####################
+
+        # S = source length 
+        # N = batch size
+        # E = embedding dimension
+        # L = target length
+
+        aux = torch.matmul(aux, self.attention_weights)
+        # Dimension S, E for two views (2,512)
+
+        # Dimension N, S, E
+        aux_t = aux.permute(0, 2, 1)
+
+        prod = torch.bmm(aux, aux_t)
+        relu_res = self.relu(prod)
+        
+        aux_sum = torch.sum(torch.reshape(relu_res, (B, V*V)).T, dim=0).unsqueeze(0)
+        final_attention_weights = torch.div(torch.reshape(relu_res, (B, V*V)).T, aux_sum.squeeze(0))
+        final_attention_weights = final_attention_weights.T
+
+        final_attention_weights = torch.reshape(final_attention_weights, (B, V, V))
+
+        final_attention_weights = torch.sum(final_attention_weights, 1)
+
+        output = torch.mul(aux.squeeze(), final_attention_weights.unsqueeze(-1))
+
+        output = torch.sum(output, 1)
+
+        return output.squeeze(), final_attention_weights
+
+
+class ViewMaxAggregate(nn.Module):
+    def __init__(self,  model, lifting_net=nn.Sequential()):
+        super().__init__()
+        self.model = model
+        self.lifting_net = lifting_net
+
+    def forward(self, mvimages):
+        B, V, C, D, H, W = mvimages.shape # Batch, Views, Channel, Depth, Height, Width
+        aux = self.lifting_net(unbatch_tensor(self.model(batch_tensor(mvimages, dim=1, squeeze=True)), B, dim=1, unsqueeze=True))
+        pooled_view = torch.max(aux, dim=1)[0]
+        return pooled_view.squeeze(), aux
+
+
+class ViewAvgAggregate(nn.Module):
+    def __init__(self,  model, lifting_net=nn.Sequential()):
+        super().__init__()
+        self.model = model
+        self.lifting_net = lifting_net
+
+    def forward(self, mvimages):
+        B, V, C, D, H, W = mvimages.shape # Batch, Views, Channel, Depth, Height, Width
+        aux = self.lifting_net(unbatch_tensor(self.model(batch_tensor(mvimages, dim=1, squeeze=True)), B, dim=1, unsqueeze=True))
+        pooled_view = torch.mean(aux, dim=1)
+        return pooled_view.squeeze(), aux
+
+class TemporalAggregation(nn.Module):
+    def __init__(self, temporal_type, hidden_dim, window_size, dropout=0.1,
+                 use_position_encoding=False, num_attention_heads=4, lstm_dropout=0.3):
+        super().__init__()
+
+        self.num_attention_heads = num_attention_heads
+        self.temporal_type = temporal_type
+        self.hidden_dim = hidden_dim
+        self.feat_dim = hidden_dim * 2 if temporal_type == "bilstm" else hidden_dim
+        self.use_position_encoding = use_position_encoding
+        self.lstm_dropout = lstm_dropout
+
+        # learnable temporal position encoding (only used when explicitly enabled)
+        if self.use_position_encoding:
+            self.temporal_position_encoding = nn.Parameter(
+                torch.randn(1, window_size, hidden_dim) * 0.02
+            )
+
+        # build temporal module
+        self.temporal = self._build_temporal_module(temporal_type, hidden_dim, dropout)
+
+    def _build_temporal_module(self, temporal_type, hidden_dim, dropout):
+        if temporal_type == 'bilstm':
+            return nn.LSTM(
+                hidden_dim, hidden_dim, num_layers=2,
+                batch_first=True, bidirectional=True, dropout=self.lstm_dropout
+            )
+
+        elif temporal_type == 'tcn':
+            return nn.Sequential(
+                nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1),
+                nn.ReLU(),
+                nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1)
+            )
+
+        elif temporal_type == 'attention':
+            return nn.MultiheadAttention(
+                hidden_dim, num_heads=self.num_attention_heads, 
+                dropout=dropout, batch_first=True
+            )
+
+        else:  # avgpool, maxpool
+            return None
+
+    def forward(self, x):
+        seq_len = x.size(1)
+
+        if self.use_position_encoding:
+            x = x + self.temporal_position_encoding[:, :seq_len, :]
+
+        if self.temporal_type == 'avgpool':
+            x = torch.mean(x, dim=1)
+
+        elif self.temporal_type == 'maxpool':
+            x = torch.max(x, dim=1)[0]
+
+        elif self.temporal_type == 'tcn':
+            x = x.permute(0, 2, 1)
+            x = self.temporal(x)
+            x = x.permute(0, 2, 1)
+            x = torch.max(x, dim=1)[0]
+
+        elif self.temporal_type == 'attention':
+            x, _ = self.temporal(x, x, x)
+            x = torch.max(x, dim=1)[0]
+
+        elif self.temporal_type == 'bilstm':
+            lstm_out, _ = self.temporal(x)
+            x = torch.max(lstm_out, dim=1)[0]
+
+        return x

opensportslib/models/utils/common.py

@@ -0,0 +1,176 @@
+"""
+Copyright 2022 James Hong, Haotian Zhang, Matthew Fisher, Michael Gharbi,
+Kayvon Fatahalian
+
+Redistribution and use in source and binary forms, with or without modification,
+are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+this list of conditions and the following disclaimer in the documentation and/or
+other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors
+may be used to endorse or promote products derived from this software without
+specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
+ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+"""
+import abc
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ABCModel:
+
+    @abc.abstractmethod
+    def get_optimizer(self, opt_args):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def epoch(self, loader, **kwargs):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def predict(self, seq):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def state_dict(self):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def load(self, state_dict):
+        raise NotImplementedError()
+
+
+class BaseRGBModel(ABCModel):
+
+    def get_optimizer(self, opt_args):
+        return torch.optim.AdamW(self._get_params(), **opt_args), (
+            torch.cuda.amp.GradScaler() if self.device == "cuda" else None
+        )
+
+    """ Assume there is a self._model """
+
+    def _get_params(self):
+        return list(self._model.parameters())
+
+    def state_dict(self):
+        if isinstance(self._model, nn.DataParallel):
+            return self._model.module.state_dict()
+        return self._model.state_dict()
+
+    def load(self, state_dict):
+        if isinstance(self._model, nn.DataParallel):
+            self._model.module.load_state_dict(state_dict)
+        else:
+            self._model.load_state_dict(state_dict)
+
+
+def step(optimizer, scaler, loss, lr_scheduler=None, backward_only=False):
+    """
+    Perform a backward pass, and optionally update the model parameters and learning rate scheduler.
+
+    Args:
+        optimizer (torch.optim.Optimizer): The optimizer to update model parameters.
+        scaler (torch.cuda.amp.GradScaler): The gradient scaler for mixed precision training.
+        loss (torch.Tensor): The computed loss to backpropagate.
+        lr_scheduler : The learning rate scheduler.
+            Default: None.
+        backward_only (bool): If True, only perform the backward pass.
+            Default: False.
+    """
+    if scaler is None:
+        loss.backward()
+    else:
+        scaler.scale(loss).backward()
+
+    if not backward_only:
+        if scaler is None:
+            optimizer.step()
+        else:
+            scaler.step(optimizer)
+            scaler.update()
+        if lr_scheduler is not None:
+            lr_scheduler.step()
+        optimizer.zero_grad()
+
+
+class SingleStageGRU(nn.Module):
+
+    def __init__(self, in_dim, hidden_dim, out_dim, num_layers=5):
+        super(SingleStageGRU, self).__init__()
+        self.backbone = nn.GRU(
+            in_dim,
+            hidden_dim,
+            num_layers=num_layers,
+            batch_first=True,
+            bidirectional=True,
+        )
+        self.fc_out = nn.Sequential(
+            nn.BatchNorm1d(2 * hidden_dim),
+            nn.Dropout(),
+            nn.Linear(2 * hidden_dim, out_dim),
+        )
+
+    def forward(self, x):
+        batch_size, clip_len, _ = x.shape
+        x, _ = self.backbone(x)
+        x = self.fc_out(x.reshape(-1, x.shape[-1]))
+        return x.view(batch_size, clip_len, -1)
+
+
+class SingleStageTCN(nn.Module):
+
+    class DilatedResidualLayer(nn.Module):
+        def __init__(self, dilation, in_channels, out_channels):
+            super(SingleStageTCN.DilatedResidualLayer, self).__init__()
+            self.conv_dilated = nn.Conv1d(
+                in_channels, out_channels, 3, padding=dilation, dilation=dilation
+            )
+            self.conv_1x1 = nn.Conv1d(out_channels, out_channels, 1)
+            self.dropout = nn.Dropout()
+
+        def forward(self, x, mask):
+            out = F.relu(self.conv_dilated(x))
+            out = self.conv_1x1(out)
+            out = self.dropout(out)
+            return (x + out) * mask[:, 0:1, :]
+
+    def __init__(self, in_dim, hidden_dim, out_dim, num_layers, dilate):
+        super(SingleStageTCN, self).__init__()
+        self.conv_1x1 = nn.Conv1d(in_dim, hidden_dim, 1)
+        self.layers = nn.ModuleList(
+            [
+                SingleStageTCN.DilatedResidualLayer(
+                    2**i if dilate else 1, hidden_dim, hidden_dim
+                )
+                for i in range(num_layers)
+            ]
+        )
+        self.conv_out = nn.Conv1d(hidden_dim, out_dim, 1)
+
+    def forward(self, x, m=None):
+        batch_size, clip_len, _ = x.shape
+        if m is None:
+            m = torch.ones((batch_size, 1, clip_len), device=x.device)
+        else:
+            m = m.permute(0, 2, 1)
+        x = self.conv_1x1(x.permute(0, 2, 1))
+        for layer in self.layers:
+            x = layer(x, m)
+        x = self.conv_out(x) * m[:, 0:1, :]
+        return x.permute(0, 2, 1)