-
Notifications
You must be signed in to change notification settings - Fork 0
/
unet.py
183 lines (143 loc) · 5.82 KB
/
unet.py
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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
""" Parts of the U-Net model """
#%%
import torch
from torch import cuda
from torch._C import Value, device
from torch.distributed import is_available
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.modules.upsampling import Upsample
#Encoder:使得模型理解了图像的内容,但是丢弃了图像的位置信息。
#Decoder:使模型结合Encoder对图像内容的理解,恢复图像的位置信息。
class DoubleConv(nn.Module):
"""convolution->BN->Relu"""
def __init__(self,in_channels,out_channels,mid_channels = None):
super().__init__()
#这部分和论文不一样,多了个中间输入的channel
#如果使用中间channel,然后在卷积到out_channel
if not mid_channels:
mid_channels = out_channels
self.double_conv = nn.Sequential(
nn.Conv2d(in_channels,mid_channels,kernel_size=3,padding=1),
nn.BatchNorm2d(mid_channels),
nn.ReLU(),
nn.Conv2d(mid_channels,out_channels,kernel_size=3,padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(),
)
def forward(self,x):
return self.double_conv(x)
#编码器encoder
class Down(nn.Module):
"""Downscaling with maxpool then double conv"""
def __init__(self, in_channels,out_channels,maxPool = True):
super().__init__()
maxPool_conv = nn.Sequential(
nn.MaxPool2d(2),
DoubleConv(in_channels,out_channels),
)
#使用卷积代替maxPoling 下采样,这样不会丢失位置信息
#但是如果网络太深,会产生过拟合
down_conv = nn.Sequential(
nn.Conv2d(in_channels,out_channels,kernel_size=3,stride=2,padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(),
DoubleConv(in_channels,out_channels),
)
self.downsample = ( maxPool_conv if maxPool else down_conv)
def forward(self,x):
return self.downsample(x)
#解码器Decoder
class Up(nn.Module):
"""Downscaling with maxpool then double conv"""
def __init__(self,in_channels,out_channels, bilinear=True):
super().__init__()
#如果是双线性插值,使用普通卷积来减少通道
if bilinear:
self.up = Upsample(scale_factor=2,mode='bilinear',align_corners=True)
self.conv = DoubleConv(in_channels,out_channels,in_channels // 2)
else:
#采用转置卷积代替上采样,out_channel 是in_channels的一半
self.up = nn.ConvTranspose2d(in_channels,in_channels//2,kernel_size=2,stride=2)
self.conv = DoubleConv(in_channels, out_channels)
def forward(self,x1,x2):
#据论文,需要进行融合还原之前的尺寸
#x1是上采样获得的特征
#x2是下采样获得的特征
x1 = self.up(x1)
if (x1.size(2) != x2.size(2)) or (x1.size(3) != x2.size(3)):
#input is CHW
#这个是解决填充不一致的问题
diffY = x2.size()[2] - x1.size()[2]
diffX = x2.size()[3] - x1.size()[3]
#print('sizes',x1.size(),x2.size(),diffX // 2, diffX - diffX//2, diffY // 2, diffY - diffY//2)
x1 = F.pad(x1, (diffX // 2, diffX - diffX//2,
diffY // 2, diffY - diffY//2))
#print("pad x1:",x1.size())
x = torch.cat([x2, x1], dim=1)
x = self.conv(x)
return x
class outlayer(nn.Module):
def __init__(self,in_channels,out_channels):
super().__init__()
self.conv = nn.Conv2d(in_channels,out_channels,kernel_size=1)
def forward(self,x):
x =self.conv(x)
return x
class UNet(nn.Module):
def __init__(self,n_classes,n_channels = 3,bilinear=True) -> None:
super().__init__()
self.n_channels = n_channels
self.n_channels = n_classes
self.bilinear = bilinear
self.start = DoubleConv(n_channels,64)
self.down1 = Down(64,128)
self.down2 = Down(128,256)
self.down3 = Down(256,512)
#self.down4 = Down(512,1024)
factor = 2 if bilinear else 1
self.down4 = Down(512, 1024 // factor)
self.up1 = Up(1024, 512 // factor, bilinear)
self.up2 = Up(512, 256 // factor, bilinear)
self.up3 = Up(256, 128 // factor, bilinear)
self.up4 = Up(128, 64, bilinear)
self.final_conv = outlayer(64, n_classes)
# self.up1 = Up(1024, 512,bilinear)
# self.up2 = Up(512, 256,bilinear)
# self.up3 = Up(256, 128 ,bilinear)
# self.up4 = Up(128, 64,bilinear)
# self.final_conv = outlayer(64, n_classes)
def forward(self,x):
x0 = self.start(x) #3-64
#print(x0.shape)
x1 = self.down1(x0)#64-128
#print(f"x1.shape:\n{x1.shape}")
x2 = self.down2(x1)#128-246
#print(f"x2.shape:\n{x2.shape}")
x3 = self.down3(x2)#256-512
#print(f"x3.shape:\n{x3.shape}")
x4 = self.down4(x3)#512-1024
#print(f"x4.shape:\n{x4.shape}")
x = self.up1(x4, x3)#1024-512
#print(f"x.shape:\n{x.shape}")
x = self.up2(x, x2)#512-256
#print(f"x.shape:\n{x.shape}")
x = self.up3(x, x1)#256-128
#print(f"x.shape:\n{x.shape}")
x = self.up4(x, x0)#128-64
#print(f"x.shape:\n{x.shape}")
logits = self.final_conv(x)
#print(f"logits:\n{logits.shape}")
return logits
if __name__ == '__main__':
net = UNet(n_channels=1,n_classes=2,bilinear=False)
dev = ('cuda:0' if torch.cuda.is_available() else 'cpu')
with open('log.txt','w') as fp:
for name,param in net.named_parameters():
print(f"name={name}:param={param}",file = fp)
# print(dev)
# x = torch.randn(1,1,572,572)
# out = net(x).to(dev)
# print(net)
# print(out.shape)
# %%