unet2.0/model.py at main · fire111345/unet2.0 · 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
108
109
110
111
import torch
import torch.nn as nn


class Unet(nn.Module):
    def __init__(self, in_channels=1, out_channels=1):
        super().__init__()

        # Encoder
        self.conv1 = nn.Sequential(
            nn.Conv3d(in_channels=in_channels, out_channels=32, kernel_size=3, padding=1),
            nn.BatchNorm3d(32),
            nn.ReLU(inplace=True),
            nn.Conv3d(in_channels=32, out_channels=32, kernel_size=3, padding=1),
            nn.BatchNorm3d(32),
            nn.ReLU(inplace=True)
        )
        self.pool1 = nn.MaxPool3d(kernel_size=2)

        self.conv2 = nn.Sequential(
            nn.Conv3d(in_channels=32, out_channels=64, kernel_size=3, padding=1),
            nn.BatchNorm3d(64),
            nn.ReLU(inplace=True),
            nn.Conv3d(in_channels=64, out_channels=64, kernel_size=3, padding=1),
            nn.BatchNorm3d(64),
            nn.ReLU(inplace=True)
        )
        self.pool2 = nn.MaxPool3d(kernel_size=2)

        # Decoder
        self.up2 = nn.ConvTranspose3d(
            in_channels=64, out_channels=64,
            kernel_size=2, stride=2
        )
        self.dec2 = nn.Sequential(
            nn.Conv3d(in_channels=64 + 64, out_channels=64, kernel_size=3, padding=1),
            nn.BatchNorm3d(64),
            nn.ReLU(inplace=True),
            nn.Conv3d(in_channels=64, out_channels=64, kernel_size=3, padding=1),
            nn.BatchNorm3d(64),
            nn.ReLU(inplace=True)
        )

        self.up1 = nn.ConvTranspose3d(
            in_channels=64, out_channels=32,
            kernel_size=2, stride=2
        )
        self.dec1 = nn.Sequential(
            nn.Conv3d(in_channels=32 + 32, out_channels=32, kernel_size=3, padding=1),
            nn.BatchNorm3d(32),
            nn.ReLU(inplace=True),
            nn.Conv3d(in_channels=32, out_channels=32, kernel_size=3, padding=1),
            nn.BatchNorm3d(32),
            nn.ReLU(inplace=True)
        )

        self.final = nn.Conv3d(32, out_channels, kernel_size=1)

    def forward(self, x):
        # Encoder
        s1 = self.conv1(x)
        #print(s1.shape)

        x = self.pool1(s1)
        #print(x.shape)

        s2 = self.conv2(x)
        #print(s2.shape)

        x = self.pool2(s2)
        #print(x.shape)

        # Decoder
        x = self.up2(x)
        #print(x.shape)

        x = torch.cat([x, s2], dim=1)
        #print(x.shape)

        x = self.dec2(x)
        #print(x.shape)

        x = self.up1(x)
        #print(x.shape)

        x = torch.cat([x, s1], dim=1)
        #print(x.shape)

        x = self.dec1(x)
        #print(x.shape)

        return  self.final(x)

# torch.Size([4, 32, 16, 256, 256])
# torch.Size([4, 32, 8, 128, 128])
# torch.Size([4, 64, 8, 128, 128])
# torch.Size([4, 64, 4, 64, 64])
# torch.Size([4, 64, 8, 128, 128])
# torch.Size([4, 128, 8, 128, 128])
# torch.Size([4, 64, 8, 128, 128])
# torch.Size([4, 32, 16, 256, 256])
# torch.Size([4, 64, 16, 256, 256])
# torch.Size([4, 32, 16, 256, 256])
# out unet shape is torch.Size([4, 32, 16, 256, 256])

#   net = Unet(in_channels=1, out_channels=4)
# if __name__ == "__main__":
#     unet = Unet(in_channels=1, out_channels=3)
#     test = torch.rand((1, , 16, 256, 256))  # B C D H W
#     out = unet(test)
#     print("out unet shape is", out.shape)