NeuralDecodingVideo/model.py at main · ganglab/NeuralDecodingVideo · 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
import torch
import torch.nn as nn
import torch.nn.functional as F

class SpikeDecoder(nn.Module):

    def __init__(self, ncell):

        super(SpikeDecoder, self).__init__()

        self.dense1 = nn.Linear(ncell, 512)
        self.bn1 = nn.BatchNorm2d(45)
        self.dense2 = nn.Linear(512, 1024*3)
        self.reshape = nn.Unflatten(2, (3 ,32, 32))

    def forward(self, x):

        x = self.dense1(x)
        x= x.unsqueeze(-1)
        x = self.bn1(x)
        x = x.squeeze(-1)
        x = F.relu(x)
        x = torch.sigmoid(self.dense2(x))
        x = self.reshape(x)
        x=x.permute(0,2,1,3,4)

        return x

class VideoUNet(nn.Module):

    def __init__(self, input_shape, ncell,hiddensize=64):

        super(VideoUNet, self).__init__()

        self.Dense= SpikeDecoder(ncell)
        self.encoder = nn.Sequential(
            nn.Conv3d(input_shape[0], hiddensize, kernel_size=(7, 7, 7),padding='same'),
            nn.BatchNorm3d(hiddensize),
            nn.ReLU(),
            nn.MaxPool3d([3,2,2]),
            nn.Conv3d(hiddensize, hiddensize*2, kernel_size=(5,5, 5), padding='same'),
            nn.BatchNorm3d(hiddensize*2),
            nn.ReLU(),
            nn.MaxPool3d([3,2,2]),
            nn.Conv3d(hiddensize*2, hiddensize*4, kernel_size=(3, 3, 3),  padding='same'),
            nn.BatchNorm3d(hiddensize*4),
            nn.ReLU(),
            nn.MaxPool3d([5,2,2]),
            nn.Conv3d(hiddensize*4, hiddensize*4, kernel_size=(3, 3, 3),  padding='same'),
            nn.BatchNorm3d(hiddensize*4),
            nn.ReLU(),
            nn.MaxPool3d([1,2,2]),
        )

        self.decoder = nn.Sequential(
            nn.Upsample(scale_factor=tuple([1,2,2]), mode='nearest'),
            nn.Conv3d(hiddensize*4, hiddensize*4, kernel_size=(3, 3, 3), padding='same'),
            nn.BatchNorm3d(hiddensize*4),
            nn.ReLU(),
            nn.Upsample(scale_factor=tuple([5,2,2]), mode='nearest'),
            nn.Conv3d(hiddensize*4,hiddensize*2, kernel_size=(3, 3, 3), padding='same'),
            nn.BatchNorm3d(hiddensize*2),
            nn.ReLU(),
            nn.Upsample(scale_factor=tuple([3,2,2]), mode='nearest'),
            nn.Conv3d(hiddensize*2, hiddensize, kernel_size=(5,5,5), padding='same'),
            nn.BatchNorm3d(hiddensize),
            nn.ReLU(),
            nn.Upsample(scale_factor=tuple([3,2,2]), mode='nearest'),
            nn.Conv3d(hiddensize, 3, kernel_size=(7, 7, 7), padding='same'),
            nn.BatchNorm3d(3)
        )

        self.tanh = nn.Tanh()

    def forward(self, x):

        x=self.Dense(x)
        x = self.encoder(x)
        x = self.decoder(x)
        x=(self.tanh(x)+1)/2

        return x