nn_estimation/models.py at main · kylejray/nn_estimation · 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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
from training import ModelBuilder
import torch
from torch import nn


# simple network for rate based estimation
# also serves as an example for making new models
class SingleTimeStep(ModelBuilder):
    def __init__(self, options = None):
        default_options = {'n_input':2, 'n_output':1, 'n_hidden':512, 'num_inner':2}
        super().__init__(default_options, options)

    def generate_network(self):
        opt = self.options
        self.h = fully_connected_linear(
            opt.n_input, opt.n_output, opt.n_hidden, opt.num_inner,
            )
        return

    def forward(self, s):
        return self.h(s)

class SingleTimeStep_Tanh(ModelBuilder):
    def __init__(self, options = None):
        default_options = {'n_input':2, 'n_output':1, 'n_hidden':512, 'num_inner':2}
        super().__init__(default_options, options)

    def generate_network(self):
        opt = self.options
        self.h = fully_connected_linear_tanh(
            opt.n_input, opt.n_output, opt.n_hidden, opt.num_inner,
            )
        return

    def forward(self, s):
        return self.h(s)

# here is using two linear networks to also collapse all trajectories to one number
# potentially for testing the TUT or the oneshot method
class FullTrajectory(ModelBuilder):
    def __init__(self, options = None):
        default_options = {'n_input':2, 'n_output':1, 'n_hidden':512, 'num_inner':2, 'traj_length':10}
        super().__init__(default_options, options)

    def generate_network(self):
        opt = self.options

        self.h0 = fully_connected_linear(
            opt.n_input, opt.n_output, opt.n_hidden, opt.num_inner,
            )

        self.h1 = fully_connected_linear(
            opt.traj_length, opt.n_output, opt.n_hidden, opt.num_inner,
            )
        return

    def forward(self, s):
        w_coords = self.h0(s).squeeze()
        w = self.h1(w_coords)
        return h1


def fully_connected_linear(n_input, n_output, n_hidden, num_inner):
    '''
    helper function to make fully connected linear/relu layers
    '''
    linear_layers = [nn.Linear(n_input, n_hidden),
                    *[nn.Linear(n_hidden, n_hidden) for i in range(num_inner)],
                    nn.Linear(n_hidden, n_output),
                    ]
    relu_layers = [nn.ReLU(inplace=True) for i in range(len(linear_layers)-1)]
    layers = [None]*(2*len(linear_layers)-1)
    layers[::2] = linear_layers
    layers[1::2] = relu_layers
    return nn.Sequential(*layers)

def fully_connected_linear_tanh(n_input, n_output, n_hidden, num_inner):
    '''
    helper function to make fully connected linear/relu layers
    '''
    linear_layers = [nn.Linear(n_input, n_hidden),
                    *[nn.Linear(n_hidden, n_hidden) for i in range(num_inner)],
                    nn.Linear(n_hidden, n_output),
                    ]
    tanh_layers = [nn.Tanh() for i in range(len(linear_layers)-1)]
    layers = [None]*(2*len(linear_layers)-1)
    layers[::2] = linear_layers
    layers[1::2] = tanh_layers
    return nn.Sequential(*layers)

class TimeOddCurrent(FullTrajectory):
    '''
    return a time odd function of a full trajectory
    1D only right now
    '''

    def __init__(self, options):
        super().__init__(options)

    def one_pass(self, s):
        h0 = self.h0(s).squeeze()
        return self.h1(h0)

    def forward(self, s):
        s_rev = torch.fliplr(s)
        s_rev[...,1] *= -1
        return self.one_pass(s) - self.one_pass(s_rev)