GraphTransformer_BrainAge/ConvNet.py at main · anaharrismatnez/GraphTransformer_BrainAge · 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
import torch.nn as nn
import torch.nn.functional as F


class ConvBlock(nn.Module):
    def __init__(self, inChannels, outChannels):
        super(ConvBlock, self).__init__()
        self.conv1 = nn.Conv3d(inChannels, outChannels, kernel_size=3, padding=1, stride=1)
        self.bn1 = nn.BatchNorm3d(outChannels)
        self.conv2 = nn.Conv3d(outChannels, outChannels, kernel_size=3, padding=1, stride=1)
        self.bn2 = nn.BatchNorm3d(outChannels)

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = F.relu(self.bn2(self.conv2(out)))
        return out


class DeConvLayer(nn.Module):
    def __init__(self, inChannels, outChannels):
        super(DeConvLayer, self).__init__()
        self.deconv = nn.ConvTranspose3d(inChannels, outChannels, kernel_size=2, stride=2)
        self.bn = nn.BatchNorm3d(outChannels)

    def forward(self, x):
        out = F.relu(self.bn(self.deconv(x)))
        return out


class ConvAutoEncoder(nn.Module):
    def __init__(self, nChannels):
        super(ConvAutoEncoder, self).__init__()
        # encoder
        self.Conv1 = ConvBlock(1, nChannels)
        self.Conv2 = ConvBlock(nChannels, nChannels*2)
        self.Conv3 = ConvBlock(nChannels*2, nChannels*4)
        self.Conv4 = ConvBlock(nChannels*4, nChannels*8)
        self.AvgPool = nn.AvgPool3d(kernel_size=2, stride=2)

        # decoder
        self.deConv1 = DeConvLayer(nChannels*8, nChannels*4)
        self.Conv5 = ConvBlock(nChannels*4, nChannels*4)
        self.deConv2 = DeConvLayer(nChannels*4, nChannels*2)
        self.Conv6 = ConvBlock(nChannels*2, nChannels*2)
        self.deConv3 = DeConvLayer(nChannels*2, nChannels)
        self.Conv7 = ConvBlock(nChannels, nChannels)
        self.Conv8 = nn.Conv3d(nChannels, 1, kernel_size=1)

    def forward(self, x):
        feature = self.AvgPool(self.Conv1(x))
        feature = self.AvgPool(self.Conv2(feature))
        feature = self.AvgPool(self.Conv3(feature))
        feature = self.deConv1(self.Conv4(feature))
        feature = self.deConv2(self.Conv5(feature))
        feature = self.deConv3(self.Conv6(feature))
        out = self.Conv8(self.Conv7(feature))
        return out


class Feature_Extraction(nn.Module):
    def __init__(self, nChannels):
        super(Feature_Extraction, self).__init__()
        # encoder for feature extraction
        self.Conv1 = ConvBlock(1, nChannels)
        self.Conv2 = ConvBlock(nChannels, nChannels*2)
        self.Conv3 = ConvBlock(nChannels*2, nChannels*4)
        self.Conv4 = ConvBlock(nChannels*4, nChannels*8)
        self.AvgPool = nn.AvgPool3d(kernel_size=2, stride=2)

    def forward(self, x):
        feature = self.AvgPool(self.Conv1(x))
        feature = self.AvgPool(self.Conv2(feature))
        out_m = self.Conv3(feature)
        feature = self.AvgPool(out_m)
        out_s = self.Conv4(feature)
        return out_s, out_m