Upload New File

models/ae_3dconv.py

+# baseline AE models for video data
+from __future__ import absolute_import, print_function
+import torch
+from torch import nn
+
+class AutoEncoderCov3D(nn.Module):
+    def __init__(self, chnum_in):
+        super(AutoEncoderCov3D, self).__init__()
+        self.chnum_in = chnum_in # input channel number is 1;
+        feature_num = 128
+        feature_num_2 = 96
+        feature_num_x2 = 256
+        self.encoder = nn.Sequential(
+            nn.Conv3d(self.chnum_in, feature_num_2, (3,3,3), stride=(1, 2, 2), padding=(1, 1, 1),),
+            nn.BatchNorm3d(feature_num_2),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv3d(feature_num_2, feature_num, (3,3,3), stride=(2,2,2), padding=(1,1,1)),
+            nn.BatchNorm3d(feature_num),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv3d(feature_num, feature_num_x2, (3,3,3), stride=(2,2,2), padding=(1,1,1)),
+            nn.BatchNorm3d(feature_num_x2),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv3d(feature_num_x2, feature_num_x2, (3,3,3), stride=(2, 2, 2), padding=(1, 1, 1)),
+            nn.BatchNorm3d(feature_num_x2),
+            nn.LeakyReLU(0.2, inplace=True)
+        )
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose3d(feature_num_x2, feature_num_x2, (3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1), output_padding=(1, 1, 1)),
+            nn.BatchNorm3d(feature_num_x2),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.ConvTranspose3d(feature_num_x2, feature_num, (3,3,3), stride=(2,2,2), padding=(1,1,1), output_padding=(1,1,1)),
+            nn.BatchNorm3d(feature_num),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.ConvTranspose3d(feature_num, feature_num_2, (3,3,3), stride=(2,2,2), padding=(1,1,1), output_padding=(1,1,1)),
+            nn.BatchNorm3d(feature_num_2),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.ConvTranspose3d(feature_num_2, self.chnum_in, (3,3,3), stride=(1,2,2), padding=(1,1,1), output_padding=(0,1,1))
+        )
+
+    def forward(self, x):
+        f = self.encoder(x)
+        out =f
+        #out = self.decoder(f)
+        return out
+
+