-
Notifications
You must be signed in to change notification settings - Fork 2
/
PLS_buildingblocks.py
75 lines (62 loc) · 2.83 KB
/
PLS_buildingblocks.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
import torch.nn as nn
import torch
import torch.utils.checkpoint as cp
class DSConv3D(nn.Module):
def __init__(self, in_chans, out_chans, dilation=1, dstride=2, padding=1):
super(DSConv3D, self).__init__()
self.dConv = nn.Conv3d(in_chans, in_chans, kernel_size=3, stride=dstride, padding=padding,
dilation=dilation, groups=in_chans, bias=False)
self.conv = nn.Conv3d(in_chans, out_chans, kernel_size=1, dilation=1, stride=1, bias=False)
self.norm = nn.BatchNorm3d(out_chans, momentum=0.01)
self.relu = nn.ReLU(inplace=True)
def forward(self, x):
out = self.dConv(x)
out = self.conv(out)
out = self.relu(out)
return out
class DrdbBlock3D(nn.Module):
def __init__(self, in_chans, out_chans, growth_rate, nr_blocks=4):
super(DrdbBlock3D, self).__init__()
self.nr_blocks = nr_blocks
self.in_chans = in_chans
self.out_chans = out_chans
self.growth_rate = growth_rate
self.memory_efficient = True
self.ds_conv_1 = DSConv3D(in_chans=self.in_chans, out_chans=growth_rate, dilation=1, dstride=1,
padding=1)
self.ds_conv_2 = DSConv3D(in_chans=self.in_chans + growth_rate, out_chans=growth_rate, dilation=2,
dstride=1, padding=2)
self.ds_conv_3 = DSConv3D(in_chans=self.in_chans + growth_rate * 2, out_chans=growth_rate,
dilation=3, dstride=1, padding=3)
self.ds_conv_4 = DSConv3D(in_chans=self.in_chans + growth_rate * 3, out_chans=growth_rate,
dilation=4, dstride=1, padding=4)
self.conv = nn.Conv3d(in_chans + growth_rate * 4, self.out_chans, kernel_size=1)
def forward(self, x):
if self.memory_efficient:
out = cp.checkpoint(self.bottleneck_function, x)
else:
out = self.bottleneck_function(x)
return out
def bottleneck_function(self, x):
out = self.ds_conv_1(x)
cat = torch.cat([out, x], 1)
out = self.ds_conv_2(cat)
cat = torch.cat([out, cat], 1)
out = self.ds_conv_3(cat)
cat = torch.cat([out, cat], 1)
out = self.ds_conv_4(cat)
cat = torch.cat([out, cat], 1)
out = self.conv(cat)
out = torch.add(out, x)
return out
class DecoderBlock(nn.Module):
def __init__(self, in_chans, out_chans):
super(DecoderBlock, self).__init__()
self.in_chans = in_chans
self.out_chans = out_chans
self.ds_conv = DSConv3D(in_chans=in_chans, out_chans=out_chans, dilation=1, dstride=1)
self.upsampled = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=False)
def forward(self, x):
out = self.ds_conv(x)
out = self.upsampled(out)
return out