Global loss and shared individually connected layers about gla-gcn HOT 4 OPEN

class S_AGCN(TemporalModelBase):

      def __init__(self, num_joints_in, in_features, num_joints_out,
                   filter_widths, causal=False, dropout=0.25, channels=96, dataset='h36m'):
  
          super().__init__(num_joints_in, in_features, num_joints_out, filter_widths, causal, dropout, channels, dataset)
  
          layers_tcngcn = []
          num_person = 1
          in_channels = 2
          num_point = num_joints_in
          self.graph = Graph(dataset)
          self.data_bn = nn.BatchNorm1d(num_person * in_channels * num_point)
          A = self.graph.A
          self.expand_gcn = TCN_GCN_unit(2, channels, A)
  
          self.causal_shift = []
          next_dilation = filter_widths[0]
          for i in range(0, len(filter_widths)):
              self.pad.append((filter_widths[i] - 1)*next_dilation // 2)
              self.causal_shift.append((filter_widths[i]//2 * next_dilation) if causal else 0)
              layers_tcngcn.append(TCN_GCN_unit(channels, channels, A))
              layers_tcngcn.append(TCN_GCN_unit(channels, channels, A, stride=filter_widths[i], residual=False))
              next_dilation = next_dilation * filter_widths[i]
  
          self.layers_tcngcn = nn.ModuleList(layers_tcngcn)
          self.fc = nn.Conv1d(channels, 3, 1)
  
      def set_bn_momentum(self, momentum):
          self.data_bn.momentum = momentum
          self.expand_gcn.gcn1.bn.momentum = momentum
          self.expand_gcn.tcn1.bn.momentum = momentum
          for layer in self.layers_tcngcn:
              layer.gcn1.bn.momentum = momentum
              layer.tcn1.bn.momentum = momentum
  
      def _forward_blocks(self, x):
  
          N, V, T = x.size()
          v = torch.div(V, 2, rounding_mode='floor') # number of 2D pose joints
          x = self.data_bn(x)
          x = x.view(N, 1, v, 2, T)
          x = x.permute(0, 1, 3, 4, 2).contiguous()
          x = x.view(N, 2, T, v)
  
          x = self.expand_gcn(x)
          for i in range( len(self.pad) -1):
              res = x[:, :, self.causal_shift[i] + self.filter_widths[i]//2 :: self.filter_widths[i], :]
  
              x = self.drop(self.layers_tcngcn[2*i](x))
              x = self.drop(self.layers_tcngcn[2*i+1](x))
              x = res + x
          pose_3d_ = x
  
          pose_3d = torch.zeros(N, 3, v).cuda(0)
          for i in range(0,v):
              pose_joint_3d = pose_3d_[:,:,:,i].mean(2)
              pose_joint_3d = pose_joint_3d.view(N, -1, 1)
              pose_joint_3d = self.fc(pose_joint_3d)
              pose_3d[:,:,i] = pose_joint_3d.view(N,-1)
  
          return pose_3d