opensportslib 0.0.1.dev2__py3-none-any.whl

opensportslib/models/utils/impl/asformer.py

@@ -0,0 +1,390 @@
+"""
+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.
+"""
+# Adapted from: https://github.com/ChinaYi/ASFormer
+# Source license:
+"""
+MIT License
+
+Copyright (c) 2021 ChinaYi
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import copy
+import numpy as np
+import math
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def exponential_descrease(idx_decoder, p=3):
+    return math.exp(-p*idx_decoder)
+
+
+class AttentionHelper(nn.Module):
+    def __init__(self):
+        super(AttentionHelper, self).__init__()
+        self.softmax = nn.Softmax(dim=-1)
+
+
+    def scalar_dot_att(self, proj_query, proj_key, proj_val, padding_mask):
+        '''
+        scalar dot attention.
+        :param proj_query: shape of (B, C, L) => (Batch_Size, Feature_Dimension, Length)
+        :param proj_key: shape of (B, C, L)
+        :param proj_val: shape of (B, C, L)
+        :param padding_mask: shape of (B, C, L)
+        :return: attention value of shape (B, C, L)
+        '''
+        m, c1, l1 = proj_query.shape
+        m, c2, l2 = proj_key.shape
+
+        assert c1 == c2
+
+        energy = torch.bmm(proj_query.permute(0, 2, 1), proj_key)  # out of shape (B, L1, L2)
+        attention = energy / np.sqrt(c1)
+        attention = attention + torch.log(padding_mask + 1e-6) # mask the zero paddings. log(1e-6) for zero paddings
+        attention = self.softmax(attention)
+        attention = attention * padding_mask
+        attention = attention.permute(0,2,1)
+        out = torch.bmm(proj_val, attention)
+        return out, attention
+
+class AttLayer(nn.Module):
+    def __init__(self, q_dim, k_dim, v_dim, r1, r2, r3, bl, stage, att_type): # r1 = r2
+        super(AttLayer, self).__init__()
+
+        self.query_conv = nn.Conv1d(in_channels=q_dim, out_channels=q_dim // r1, kernel_size=1)
+        self.key_conv = nn.Conv1d(in_channels=k_dim, out_channels=k_dim // r2, kernel_size=1)
+        self.value_conv = nn.Conv1d(in_channels=v_dim, out_channels=v_dim // r3, kernel_size=1)
+
+        self.conv_out = nn.Conv1d(in_channels=v_dim // r3, out_channels=v_dim, kernel_size=1)
+
+        self.bl = bl
+        self.stage = stage
+        self.att_type = att_type
+        assert self.att_type in ['normal_att', 'block_att', 'sliding_att']
+        assert self.stage in ['encoder','decoder']
+
+        self.att_helper = AttentionHelper()
+        self.window_mask = self.construct_window_mask()
+
+
+    def construct_window_mask(self):
+        '''
+            construct window mask of shape (1, l, l + l//2 + l//2), used for sliding window self attention
+        '''
+        window_mask = torch.zeros((1, self.bl, self.bl + 2* (self.bl //2)))
+        for i in range(self.bl):
+            window_mask[:, i, i:i+self.bl] = 1
+        return window_mask.to(device)
+
+    def forward(self, x1, x2, mask):
+        # x1 from the encoder
+        # x2 from the decoder
+
+        query = self.query_conv(x1)
+        key = self.key_conv(x1)
+
+        if self.stage == 'decoder':
+            assert x2 is not None
+            value = self.value_conv(x2)
+        else:
+            value = self.value_conv(x1)
+
+        if self.att_type == 'normal_att':
+            return self._normal_self_att(query, key, value, mask)
+        elif self.att_type == 'block_att':
+            return self._block_wise_self_att(query, key, value, mask)
+        elif self.att_type == 'sliding_att':
+            return self._sliding_window_self_att(query, key, value, mask)
+
+            # # NOTE: verify that results are consistent regardless of batchsize
+            # tmp = self._sliding_window_self_att(query, key, value, mask)
+            # tmp2 = []
+            # for i in range(query.shape[0]):
+            #     tmp2.append(
+            #         self._sliding_window_self_att(
+            #             query[[i]], key[[i]], value[[i]], mask[[i]]))
+            # tmp2 = torch.cat(tmp2)
+            # print(tmp.shape, (tmp - tmp2).abs().mean().item())
+            # return tmp
+
+    def _normal_self_att(self,q,k,v, mask):
+        m_batchsize, c1, L = q.size()
+        _,c2,L = k.size()
+        _,c3,L = v.size()
+        padding_mask = torch.ones((m_batchsize, 1, L)).to(device) * mask[:,0:1,:]
+        output, attentions = self.att_helper.scalar_dot_att(q, k, v, padding_mask)
+        output = self.conv_out(F.relu(output))
+        output = output[:, :, 0:L]
+        return output * mask[:, 0:1, :]
+
+    def _block_wise_self_att(self, q,k,v, mask):
+        m_batchsize, c1, L = q.size()
+        _,c2,L = k.size()
+        _,c3,L = v.size()
+
+        nb = L // self.bl
+        if L % self.bl != 0:
+            q = torch.cat([q, torch.zeros((m_batchsize, c1, self.bl - L % self.bl)).to(device)], dim=-1)
+            k = torch.cat([k, torch.zeros((m_batchsize, c2, self.bl - L % self.bl)).to(device)], dim=-1)
+            v = torch.cat([v, torch.zeros((m_batchsize, c3, self.bl - L % self.bl)).to(device)], dim=-1)
+            nb += 1
+
+        padding_mask = torch.cat([torch.ones((m_batchsize, 1, L)).to(device) * mask[:,0:1,:], torch.zeros((m_batchsize, 1, self.bl * nb - L)).to(device)],dim=-1)
+
+        q = q.reshape(m_batchsize, c1, nb, self.bl).permute(0, 2, 1, 3).reshape(m_batchsize * nb, c1, self.bl)
+        padding_mask = padding_mask.reshape(m_batchsize, 1, nb, self.bl).permute(0, 2, 1, 3).reshape(m_batchsize * nb,1, self.bl)
+        k = k.reshape(m_batchsize, c2, nb, self.bl).permute(0, 2, 1, 3).reshape(m_batchsize * nb, c2, self.bl)
+        v = v.reshape(m_batchsize, c3, nb, self.bl).permute(0, 2, 1, 3).reshape(m_batchsize * nb, c3, self.bl)
+
+        output, attentions = self.att_helper.scalar_dot_att(q, k, v, padding_mask)
+        output = self.conv_out(F.relu(output))
+
+        output = output.reshape(m_batchsize, nb, c3, self.bl).permute(0, 2, 1, 3).reshape(m_batchsize, c3, nb * self.bl)
+        output = output[:, :, 0:L]
+        return output * mask[:, 0:1, :]
+
+    def _sliding_window_self_att(self, q,k,v, mask):
+        m_batchsize, c1, L = q.size()
+        _, c2, _ = k.size()
+        _, c3, _ = v.size()
+
+        # assert m_batchsize == 1  # currently, we only accept input with batch size 1
+        # padding zeros for the last segment
+        nb = L // self.bl
+        if L % self.bl != 0:
+            q = torch.cat([q, torch.zeros((m_batchsize, c1, self.bl - L % self.bl)).to(device)], dim=-1)
+            k = torch.cat([k, torch.zeros((m_batchsize, c2, self.bl - L % self.bl)).to(device)], dim=-1)
+            v = torch.cat([v, torch.zeros((m_batchsize, c3, self.bl - L % self.bl)).to(device)], dim=-1)
+            nb += 1
+        padding_mask = torch.cat([torch.ones((m_batchsize, 1, L)).to(device) * mask[:,0:1,:], torch.zeros((m_batchsize, 1, self.bl * nb - L)).to(device)],dim=-1)
+
+        # sliding window approach, by splitting query_proj and key_proj into shape (c1, l) x (c1, 2l)
+        # sliding window for query_proj: reshape
+        # NOTE: fixes issue with batchsize > 1
+        q = q.reshape(m_batchsize, c1, nb, self.bl).permute(2, 0, 1, 3).reshape(m_batchsize * nb, c1, self.bl)
+
+        # BUG: original buggy version, produces inconsistent results
+        # q = q.reshape(m_batchsize, c1, nb, self.bl).permute(0, 2, 1, 3).reshape(m_batchsize * nb, c1, self.bl)
+
+        # sliding window approach for key_proj
+        # 1. add paddings at the start and end
+        k = torch.cat([torch.zeros(m_batchsize, c2, self.bl // 2).to(device), k, torch.zeros(m_batchsize, c2, self.bl // 2).to(device)], dim=-1)
+        v = torch.cat([torch.zeros(m_batchsize, c3, self.bl // 2).to(device), v, torch.zeros(m_batchsize, c3, self.bl // 2).to(device)], dim=-1)
+        padding_mask = torch.cat([torch.zeros(m_batchsize, 1, self.bl // 2).to(device), padding_mask, torch.zeros(m_batchsize, 1, self.bl // 2).to(device)], dim=-1)
+        # 2. reshape key_proj of shape (m_batchsize*nb, c1, 2*self.bl)
+        k = torch.cat([k[:,:, i*self.bl:(i+1)*self.bl+(self.bl//2)*2] for i in range(nb)], dim=0) # special case when self.bl = 1
+        v = torch.cat([v[:,:, i*self.bl:(i+1)*self.bl+(self.bl//2)*2] for i in range(nb)], dim=0)
+        # 3. construct window mask of shape (1, l, 2l), and use it to generate final mask
+        padding_mask = torch.cat([padding_mask[:,:, i*self.bl:(i+1)*self.bl+(self.bl//2)*2] for i in range(nb)], dim=0) # of shape (m*nb, 1, 2l)
+        final_mask = self.window_mask.repeat(m_batchsize * nb, 1, 1) * padding_mask
+
+        output, attention = self.att_helper.scalar_dot_att(q, k, v, final_mask)
+        output = self.conv_out(F.relu(output))
+
+        # NOTE: fixes issue with batchsize > 1
+        output = output.reshape(nb, m_batchsize, -1, self.bl).permute(1, 2, 0, 3).reshape(m_batchsize, -1, nb * self.bl)
+
+        # BUG: original buggy version, produces inconsistent results
+        # output = output.reshape(m_batchsize, nb, -1, self.bl).permute(0, 2, 1, 3).reshape(m_batchsize, -1, nb * self.bl)
+
+        output = output[:, :, 0:L]
+        return output * mask[:, 0:1, :]
+
+
+class MultiHeadAttLayer(nn.Module):
+    def __init__(self, q_dim, k_dim, v_dim, r1, r2, r3, bl, stage, att_type, num_head):
+        super(MultiHeadAttLayer, self).__init__()
+#         assert v_dim % num_head == 0
+        self.conv_out = nn.Conv1d(v_dim * num_head, v_dim, 1)
+        self.layers = nn.ModuleList(
+            [copy.deepcopy(AttLayer(q_dim, k_dim, v_dim, r1, r2, r3, bl, stage, att_type)) for i in range(num_head)])
+        self.dropout = nn.Dropout(p=0.5)
+
+    def forward(self, x1, x2, mask):
+        out = torch.cat([layer(x1, x2, mask) for layer in self.layers], dim=1)
+        out = self.conv_out(self.dropout(out))
+        return out
+
+
+class ConvFeedForward(nn.Module):
+    def __init__(self, dilation, in_channels, out_channels):
+        super(ConvFeedForward, self).__init__()
+        self.layer = nn.Sequential(
+            nn.Conv1d(in_channels, out_channels, 3, padding=dilation, dilation=dilation),
+            nn.ReLU()
+        )
+
+    def forward(self, x):
+        return self.layer(x)
+
+
+class FCFeedForward(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(FCFeedForward, self).__init__()
+        self.layer = nn.Sequential(
+            nn.Conv1d(in_channels, out_channels, 1),  # conv1d equals fc
+            nn.ReLU(),
+            nn.Dropout(),
+            nn.Conv1d(out_channels, out_channels, 1)
+        )
+
+    def forward(self, x):
+        return self.layer(x)
+
+
+class AttModule(nn.Module):
+    def __init__(self, dilation, in_channels, out_channels, r1, r2, att_type, stage, alpha):
+        super(AttModule, self).__init__()
+        self.feed_forward = ConvFeedForward(dilation, in_channels, out_channels)
+        self.instance_norm = nn.InstanceNorm1d(in_channels, track_running_stats=False)
+        self.att_layer = AttLayer(in_channels, in_channels, out_channels, r1, r1, r2, dilation, att_type=att_type, stage=stage) # dilation
+        self.conv_1x1 = nn.Conv1d(out_channels, out_channels, 1)
+        self.dropout = nn.Dropout()
+        self.alpha = alpha
+
+    def forward(self, x, f, mask):
+        out = self.feed_forward(x)
+        out = self.alpha * self.att_layer(self.instance_norm(out), f, mask) + out
+        out = self.conv_1x1(out)
+        out = self.dropout(out)
+        return (x + out) * mask[:, 0:1, :]
+
+
+class PositionalEncoding(nn.Module):
+    "Implement the PE function."
+
+    def __init__(self, d_model, max_len=10000):
+        super(PositionalEncoding, self).__init__()
+        # Compute the positional encodings once in log space.
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) *
+                             -(math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0).permute(0,2,1) # of shape (1, d_model, l)
+        self.pe = nn.Parameter(pe, requires_grad=True)
+#         self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        return x + self.pe[:, :, 0:x.shape[2]]
+
+class Encoder(nn.Module):
+    def __init__(self, num_layers, r1, r2, num_f_maps, input_dim, num_classes, channel_masking_rate, att_type, alpha):
+        super(Encoder, self).__init__()
+        self.conv_1x1 = nn.Conv1d(input_dim, num_f_maps, 1) # fc layer
+        self.layers = nn.ModuleList(
+            [AttModule(2 ** i, num_f_maps, num_f_maps, r1, r2, att_type, 'encoder', alpha) for i in # 2**i
+             range(num_layers)])
+
+        self.conv_out = nn.Conv1d(num_f_maps, num_classes, 1)
+        self.dropout = nn.Dropout2d(p=channel_masking_rate)
+        self.channel_masking_rate = channel_masking_rate
+
+    def forward(self, x, mask):
+        '''
+        :param x: (N, C, L)
+        :param mask:
+        :return:
+        '''
+
+        if self.channel_masking_rate > 0:
+            x = x.unsqueeze(2)
+            x = self.dropout(x)
+            x = x.squeeze(2)
+
+        feature = self.conv_1x1(x)
+        for layer in self.layers:
+            feature = layer(feature, None, mask)
+
+        out = self.conv_out(feature) * mask[:, 0:1, :]
+
+        return out, feature
+
+
+class Decoder(nn.Module):
+    def __init__(self, num_layers, r1, r2, num_f_maps, input_dim, num_classes, att_type, alpha):
+        super(Decoder, self).__init__()#         self.position_en = PositionalEncoding(d_model=num_f_maps)
+        self.conv_1x1 = nn.Conv1d(input_dim, num_f_maps, 1)
+        self.layers = nn.ModuleList(
+            [AttModule(2 ** i, num_f_maps, num_f_maps, r1, r2, att_type, 'decoder', alpha) for i in # 2 ** i
+             range(num_layers)])
+        self.conv_out = nn.Conv1d(num_f_maps, num_classes, 1)
+
+    def forward(self, x, fencoder, mask):
+
+        feature = self.conv_1x1(x)
+        for layer in self.layers:
+            feature = layer(feature, fencoder, mask)
+
+        out = self.conv_out(feature) * mask[:, 0:1, :]
+
+        return out, feature
+
+class MyTransformer(nn.Module):
+    def __init__(self, num_decoders, num_layers, r1, r2, num_f_maps, input_dim, num_classes, channel_masking_rate):
+        super(MyTransformer, self).__init__()
+        self.encoder = Encoder(num_layers, r1, r2, num_f_maps, input_dim, num_classes, channel_masking_rate, att_type='sliding_att', alpha=1)
+        self.decoders = nn.ModuleList([copy.deepcopy(Decoder(num_layers, r1, r2, num_f_maps, num_classes, num_classes, att_type='sliding_att', alpha=exponential_descrease(s))) for s in range(num_decoders)]) # num_decoders
+
+
+    def forward(self, x, mask):
+        out, feature = self.encoder(x, mask)
+        outputs = out.unsqueeze(0)
+
+        for decoder in self.decoders:
+            out, feature = decoder(F.softmax(out, dim=1) * mask[:, 0:1, :], feature* mask[:, 0:1, :], mask)
+            outputs = torch.cat((outputs, out.unsqueeze(0)), dim=0)
+
+        return outputs

opensportslib/models/utils/impl/calf.py

@@ -0,0 +1,74 @@
+"""
+From "A Context-Aware Loss Function for Action Spotting in Soccer Videos"
+"""
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+class ContextAwareWeights:
+
+    def __init__(self, k1=2, k2=1, k3=2, k4=2, hit_radius=0.1, miss_radius=0.9):
+        n = k1 + k2 + k3 + k4
+        mul_w = np.ones(n)
+        add_w = np.ones(n)
+        radius = np.full(n, miss_radius)
+        for i in range(n):
+            if i < k1:
+                mul_w[i] = (k1 - i) / k1
+            elif i < k1 + k2:
+                mul_w[i] = 0
+            elif i < k1 + k2 + k3:
+                mul_w[i] = ((i - k1 - k2) - k3) / k3
+                add_w[i] = (i - k1 - k2) / k3
+                radius[i] = 1. - hit_radius
+            else:
+                mul_w[i] = (i - k1 - k2 - k3) / k4
+        self._w = np.stack([mul_w, add_w, radius], axis=1)
+        self._offset = k1 + k2
+
+    @property
+    def weights(self):
+        return self._w
+
+    @property
+    def offset(self):
+        return self._offset
+
+    def __len__(self):
+        return self._w.shape[0]
+
+
+CALF_ERROR_FLAG = True
+
+
+def set_calf_error_flag():
+    global CALF_ERROR_FLAG
+    CALF_ERROR_FLAG = 1
+
+
+def get_calf(pred, weights):
+    pred_scores = F.softmax(pred, dim=2)    # (N, L, C)
+    cl = -torch.log(
+        weights[:, :, :, 1] - pred_scores[:, :, 1:] * weights[:, :, :, 0]
+    ) + torch.log(weights[:, :, :, 2])
+    cl = torch.max(torch.zeros_like(cl), cl)
+
+    global CALF_ERROR_FLAG
+    if CALF_ERROR_FLAG:
+        tmp = torch.sum(cl)
+        if torch.isinf(tmp):
+            print('Found Inf in CALF. Supressing future errors.')
+            CALF_ERROR_FLAG = False
+        if torch.isnan(tmp):
+            print('Found NaN in CALF. Supressing future errors.')
+            CALF_ERROR_FLAG = False
+    return torch.mean(cl)
+
+
+if __name__ == '__main__':
+    c = ContextAwareWeights()
+    print(c.weights)
+    print('All 1:', -(c.weights[:, 1] - np.ones(len(c)) * c.weights[:, 0]))
+    print('All 0:', -(c.weights[:, 1] - np.zeros(len(c)) * c.weights[:, 0]))

opensportslib/models/utils/impl/gsm.py

@@ -0,0 +1,112 @@
+"""
+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.
+"""
+# Adapted PyTorch GSM implementation by Swathikiran Sudhakaran, https://github.com/swathikirans/GSM
+# Original license for GSM
+"""
+BSD 2-Clause License for GSM
+
+Copyright (c) 2019, FBK
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* 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.
+
+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
+from torch import nn
+from torch.cuda import FloatTensor as ftens
+
+
+class _GSM(nn.Module):
+    def __init__(self, fPlane, num_segments=3):
+        super(_GSM, self).__init__()
+
+        self.conv3D = nn.Conv3d(fPlane, 2, (3, 3, 3), stride=1,
+                                padding=(1, 1, 1), groups=2)
+        nn.init.constant_(self.conv3D.weight, 0)
+        nn.init.constant_(self.conv3D.bias, 0)
+        self.tanh = nn.Tanh()
+        self.fPlane = fPlane
+        self.num_segments = num_segments
+        self.bn = nn.BatchNorm3d(num_features=fPlane)
+        self.relu = nn.ReLU()
+
+    def lshift_zeroPad(self, x):
+        return torch.cat((x[:,:,1:], ftens(x.size(0), x.size(1), 1, x.size(3), x.size(4)).fill_(0)), dim=2)
+
+    def rshift_zeroPad(self, x):
+        return torch.cat((ftens(x.size(0), x.size(1), 1, x.size(3), x.size(4)).fill_(0), x[:,:,:-1]), dim=2)
+
+    def forward(self, x):
+        batchSize = x.size(0) // self.num_segments
+        shape = x.size(1), x.size(2), x.size(3)
+        assert  shape[0] == self.fPlane
+        x = x.view(batchSize, self.num_segments, *shape).permute(0, 2, 1, 3, 4).contiguous()
+        x_bn = self.bn(x)
+        x_bn_relu = self.relu(x_bn)
+        gate = self.tanh(self.conv3D(x_bn_relu))
+        gate_group1 = gate[:, 0].unsqueeze(1)
+        gate_group2 = gate[:, 1].unsqueeze(1)
+        x_group1 = x[:, :self.fPlane // 2]
+        x_group2 = x[:, self.fPlane // 2:]
+        y_group1 = gate_group1 * x_group1
+        y_group2 = gate_group2 * x_group2
+
+        r_group1 = x_group1 - y_group1
+        r_group2 = x_group2 - y_group2
+
+        y_group1 = self.lshift_zeroPad(y_group1) + r_group1
+        y_group2 = self.rshift_zeroPad(y_group2) + r_group2
+
+        y_group1 = y_group1.view(batchSize, 2, self.fPlane // 4, self.num_segments, *shape[1:]).permute(0, 2, 1, 3, 4, 5)
+        y_group2 = y_group2.view(batchSize, 2, self.fPlane // 4, self.num_segments, *shape[1:]).permute(0, 2, 1, 3, 4, 5)
+
+        y = torch.cat((y_group1.contiguous().view(batchSize, self.fPlane//2, self.num_segments, *shape[1:]),
+                       y_group2.contiguous().view(batchSize, self.fPlane//2, self.num_segments, *shape[1:])), dim=1)
+
+        return y.permute(0, 2, 1, 3, 4).contiguous().view(batchSize*self.num_segments, *shape)