FrechetMotionDistance/AutoEncoder.py at main · antmaio/FrechetMotionDistance · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
"""
This code is freely inspired by https://alanbertl.com/autoencoder-with-fast-ai/
"""

import torch.nn as nn
import torch


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

class UpSample(nn.Module):
    def __init__(self,feat_in,feat_out,out_shape=None,scale=2):
        super().__init__()
        self.conv = nn.Conv2d(feat_in,feat_out,kernel_size=(3,3),stride=1,padding=1)
        self.out_shape,self.scale = out_shape,scale


    def forward(self,x):
        return self.conv(
            nn.functional.interpolate(
                x,size=self.out_shape,scale_factor=self.scale,mode='bilinear',align_corners=True))

def get_upSamp(feat_in,feat_out, out_shape=None, scale=2, act='relu'):

    upSamp = UpSample(feat_in,feat_out,out_shape=out_shape,scale=scale).cuda()

    layer = nn.Sequential(upSamp)

    if act == 'relu':
        act_f = nn.ReLU(inplace=True).cuda()
        bn = nn.BatchNorm2d(feat_out).cuda()
        layer.add_module('ReLU',act_f)
        layer.add_module('BN',bn)
    elif act == 'sig':
        act_f = nn.Sigmoid()
        layer.add_module('Sigmoid',act_f)
    elif act == 'tanh':
        act_f = nn.Tanh()
        layer.add_module('Tanh', act_f)
    return layer

def add_layer(m,feat_in,feat_out,name,out_shape=None,scale=2,act='relu'):
    upSamp = get_upSamp(feat_in,feat_out,out_shape=out_shape,scale=scale,act=act)
    m.add_module(name,upSamp)

'''
def save_onnx(path, name, model):

    if not os.path.exists(path):
        os.makedirs(path)

    torch.onnx.export(model,
                     torch.rand((64,3,28,28)).to(device),
                     os.path.join(path, name),
                     opset_version=11,
                     do_constant_folding=True,
                     input_names = ['input'],
                     output_names = ['output'])
'''