def build_layer(f1, f2, N_layer=0, stride=1, skip=True):
"""
:param f1: 输入通道数
:param f2: 输入通道数
:param N_layer: residual blocks数量
:param skip:
"""
if stride == 1:
if f1 == f2 and skip:
layers = []
else:
# conv3x3x3 stride=1
layers = [
Conv3d_ABN(f1, f2, kernel_size=3, stride=stride, padding=1)
]
elif stride == 2:
# conv4x4x4 stride=2
layers = [Conv3d_ABN(f1, f2, kernel_size=4, stride=stride, padding=1)]
else:
layers = [Conv3d_ABN(f1, f2, kernel_size=3, stride=stride, padding=1)]
# residual bottleneck * N_layers
for i in range(N_layer):
layers.append(Bottleneck_ABN(f2))
return nn.Sequential(*layers)
def __init__(self,
n_inp=1,
feats=[32, 64, 64, 128, 128, 256, 256],
blocks=[2, 2, 2, 2, 2, 2, 2],
n_pred_p=[1, 1, 1],
n_pred_d=[4, 4, 4],
L_output=[0, 1, 2],
abn=1,
dropout_ratio=0.0):
super().__init__()
self.relu = nn.ReLU(inplace=True)
self.pixel_shuffle = nn.PixelShuffle(2)
self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
self.n_inp = n_inp
if abn == 0:
abnblock = ABN
elif abn == 1:
abnblock = InPlaceABN
else:
abnblock = InPlaceABNSync
Conv3d_ABN.define_abn(abnblock)
Bottleneck_ABN.define_abn(abnblock)
self.strides = np.array([[1, 1, 1], [2, 2, 2], [4, 4, 4], [8, 8, 8],
[16, 16, 16]])
for i in range(6):
if i == 0:
setattr(self, 'l' + str(i),
build_layer(n_inp, feats[i], blocks[i]))
setattr(self, 'r' + str(i),
build_layer(feats[i], feats[i], blocks[i], skip=True))
else:
setattr(
self, 'l' + str(i),
build_layer(feats[i - 1], feats[i], blocks[i], stride=2))
setattr(self, 'r' + str(i),
build_layer(feats[i], feats[i], blocks[i], skip=True))
setattr(self, 'rt' + str(i),
build_layer(feats[i], feats[i - 1], blocks[i]))
self.r0 = Conv3d_ABN(feats[0], feats[0], kernel_size=3, padding=1)
self.out_layer = nn.Conv3d(feats[0], 1, kernel_size=1, stride=1)
self.drop_out = nn.Dropout3d(p=dropout_ratio)
for c_out1, c_out2, out in zip(n_pred_p, n_pred_d, L_output):
setattr(self, 'out' + str(out) + 'p',
nn.Conv3d(feats[out], c_out1, kernel_size=3, padding=1))
outd = nn.Conv3d(feats[out], c_out2, kernel_size=3, padding=1)
#outd.weight.data.fill_(0.0)
outd.bias.data.fill_(0.0)
setattr(self, 'out' + str(out) + 'd', outd)
self.L_output = L_output
self.strides = self.strides[np.array(L_output)]
def build_layer(f1, f2, N_layer=0, stride=1, skip=True):
if stride == 1:
if f1 == f2 and skip:
layers = []
else:
layers = [
Conv3d_ABN(f1, f2, kernel_size=3, stride=stride, padding=1)
]
elif stride == 2:
layers = [Conv3d_ABN(f1, f2, kernel_size=4, stride=stride, padding=1)]
else:
layers = [Conv3d_ABN(f1, f2, kernel_size=3, stride=stride, padding=1)]
for i in range(N_layer):
layers.append(Bottleneck_ABN(f2))
return nn.Sequential(*layers)
def __init__(self,
n_inp=1,
feats=[32, 64, 64, 128, 128, 256, 256],
blocks=[2, 2, 2, 2, 2, 2, 2],
n_pred_p=[1, 1, 1],
n_pred_d=[4, 4, 4],
L_output=[0, 1, 2],
abn=1,
dropout_ratio=0.0):
"""
:param n_inp: 输入通道数量 [bs, n_inp, z, y, x]
:param n_pred_p: 预测类别数量 [1,1,1] 二分类
:param n_pred_d: 预测位置信息 [4,4,4] 四个坐标值
:param feats: 通道数量 [32, 64, 64, 128, 128, 256, 256]
:param blocks: residual block的数量 [2,2,2,2,2,2]
:param L_output: 输出预测结果的层 [0, 1, 2]
"""
super().__init__()
self.relu = nn.ReLU(inplace=True)
self.pixel_shuffle = nn.PixelShuffle(2)
self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
self.n_inp = n_inp
if abn == 0:
abnblock = ABN
elif abn == 1:
abnblock = InPlaceABN
else:
abnblock = InPlaceABNSync
Conv3d_ABN.define_abn(abnblock)
Bottleneck_ABN.define_abn(abnblock)
# not used
self.strides = np.array([[1, 1, 1], [2, 2, 2], [4, 4, 4], [8, 8, 8],
[16, 16, 16]])
# 6 stages
for i in range(6):
# input layer
if i == 0:
# 4 blocks
# 3x3 * 3
setattr(self, 'l' + str(i),
build_layer(n_inp, feats[i], blocks[i]))
# 3x3 * 2
setattr(self, 'r' + str(i),
build_layer(feats[i], feats[i], blocks[i], skip=True))
# other layers
else:
# 6 blocks
# stride 2 downsample
setattr(
self, 'l' + str(i),
build_layer(feats[i - 1], feats[i], blocks[i], stride=2))
setattr(self, 'r' + str(i),
build_layer(feats[i], feats[i], blocks[i], skip=True))
# add after upsampled (keep dim with previous block)
setattr(self, 'rt' + str(i),
build_layer(feats[i], feats[i - 1], blocks[i]))
self.r0 = Conv3d_ABN(feats[0], feats[0], kernel_size=3, padding=1)
self.out_layer = nn.Conv3d(feats[0], 1, kernel_size=1, stride=1)
self.drop_out = nn.Dropout3d(p=dropout_ratio)
# 不同layer的预测类别和位置信息
for c_out1, c_out2, out in zip(n_pred_p, n_pred_d, L_output):
# outp = build_layer(feats[out], feats[out] // 4, 0, skip=False)
outp = nn.Conv3d(feats[out], feats[out], kernel_size=3, padding=1)
setattr(self, 'before_out' + str(out) + 'p', outp)
# outd = build_layer(feats[out], feats[out] // 4, 0, skip=False)
outd = nn.Conv3d(feats[out], feats[out], kernel_size=3, padding=1)
setattr(self, 'before_out' + str(out) + 'd', outd)
for c_out1, c_out2, out in zip(n_pred_p, n_pred_d, L_output):
setattr(self, 'out' + str(out) + 'p',
nn.Conv3d(feats[out], c_out1, kernel_size=3, padding=1))
outd = nn.Conv3d(feats[out], c_out2, kernel_size=3, padding=1)
# outd.weight.data.fill_(0.0)
outd.bias.data.fill_(0.0)
setattr(self, 'out' + str(out) + 'd', outd)
self.L_output = L_output
# not used
self.strides = self.strides[np.array(L_output)]