text: #!/usr/bin/env python3
  # -*- coding: utf-8 -*-
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  def weights_init(m):
      if isinstance(m, nn.Conv2d):
          nn.init.kaiming_uniform_(m.weight, a=0, mode='fan_in', nonlinearity='leaky_relu')
          try:
              nn.init.constant_(m.bias, 0.01)
          except:
              pass
      if isinstance(m, nn.Linear):
          nn.init.kaiming_uniform_(m.weight, a=0, mode='fan_in', nonlinearity='leaky_relu')
          nn.init.constant_(m.bias, 0.01)
          
  class ModelGRU(nn.Module):
      def __init__(self, hidden_dim=256):
          super(ModelGRU, self).__init__()
          self.cnn_feature_dim = hidden_dim
          self.rnn_hidden_dim = hidden_dim
          self.cnn = CNN(input_dim=1, out_dim=self.cnn_feature_dim)
          self.gru = nn.GRU(
              input_size = self.cnn_feature_dim, 
              hidden_size = self.rnn_hidden_dim, 
              num_layers = 3,
              batch_first=True,
              dropout=0.2)
          self.mlp = MLP_COS(input_dim=self.rnn_hidden_dim+2)
  
      def forward(self, x, t, v0):
          batch_size, timesteps, C, H, W = x.size()
          
          x = x.view(batch_size * timesteps, C, H, W)
          x = self.cnn(x)
          
          x = x.view(batch_size, timesteps, -1)
          x, h_n, = self.gru(x)
  
          x = F.leaky_relu(x[:, -1, :])
          x = self.mlp(x, t, v0)
          return x
  
  class CNN(nn.Module):
      def __init__(self,input_dim=1, out_dim=256):
          super(CNN, self).__init__()
          self.out_dim = out_dim
          self.conv1 = nn.Conv2d(input_dim, 64, 5, stride=3, padding=2)
          self.conv2 = nn.Conv2d(64,  128, 5, stride=4, padding=2)
          self.conv3 = nn.Conv2d(128, 256, 3, stride=2, padding=1)
          self.conv4 = nn.Conv2d(256, self.out_dim, 3, stride=2, padding=1)
          
          self.bn1 = nn.BatchNorm2d(64)
          self.bn2 = nn.BatchNorm2d(128)
          self.bn3 = nn.BatchNorm2d(256)
          
          self.apply(weights_init)
  
      def forward(self, x):
          x = self.conv1(x)
          #x = self.bn1(x)
          x = F.leaky_relu(x)
          x = F.max_pool2d(x, 2, 2)
          x = self.conv2(x)
          #x = self.bn2(x)
          x = F.leaky_relu(x)
          x = F.max_pool2d(x, 2, 2)         
          x = self.conv3(x)
          #x = self.bn3(x)
          x = F.leaky_relu(x)
          x = F.max_pool2d(x, 2, 2)
          x = self.conv4(x)
          x = F.leaky_relu(x)
          x = x.view(-1, self.out_dim)
          return x
  
  class MLP_COS(nn.Module):
      def __init__(self, input_dim=257, rate=1.0, output=2):
          super(MLP_COS, self).__init__()
          self.rate = rate
          self.linear1 = nn.Linear(input_dim, 512)
          self.linear2 = nn.Linear(512, 512)
          self.linear3 = nn.Linear(512, 512)
          self.linear4 = nn.Linear(512, 256)
          self.linear5 = nn.Linear(256, output)
          
          self.apply(weights_init)
          
      def forward(self, x, t, v0):
          x = torch.cat([x, t], dim=1)
          x = torch.cat([x, v0], dim=1)
          x = self.linear1(x)
          #x = F.leaky_relu(x)
          x = torch.tanh(x)
          #x = F.dropout(x, p=0.5, training=self.training)
          x = self.linear2(x)
          #x = F.leaky_relu(x)
          x = torch.tanh(x)
          #x = F.dropout(x, p=0.5, training=self.training)
          x = self.linear3(x)
          #x = F.leaky_relu(x)
          x = torch.tanh(x)
          #x = F.dropout(x, p=0.5, training=self.training)
          
          x = self.linear4(x)
          #x = F.leaky_relu(x)
          x = torch.cos(self.rate*x)
          #x = F.dropout(x, p=0.5, training=self.training)
          x = self.linear5(x)
          return x
  