def __init__(self, num_classes, expand_ratio=4, bn_mom=0.9, bn_eps=0.001):
super().__init__()
_, _, self.conv3d = get_conv()
self._bn_mom = 1 - bn_mom
self._bn_eps = bn_eps
repeats = [3, 5, 11, 7]
strides = [
(1, 2, 2),
(1, 2, 2),
(1, 2, 2),
(1, 2, 2),
]
channels = [24, 48, 96, 192]
se = [True] * 4
input_chann = 3
stem_chann = 24
self._stem_conv_xy = self.conv3d(input_chann,
stem_chann,
stride=(1, 2, 2),
kernel_size=(1, 3, 3),
bias=False)
self._stem_conv_t = self.conv3d(stem_chann,
stem_chann,
stride=(1, 1, 1),
kernel_size=(5, 1, 1),
bias=False)
self._stem_bn = nn.BatchNorm3d(num_features=24,
momentum=self._bn_mom,
eps=self._bn_eps)
reduction = 1 / 2
reduce_pow = (reduction)**(1 / (len(repeats) - 1))
stage_expand_ratios = []
instage_expand_ratios = []
for i, block_repeats in enumerate(repeats):
ratio = reduce_pow**i
stage_ratio = expand_ratio * ratio
stage_expand_ratios.append(stage_ratio)
max_reduction = reduction * (1 / reduce_pow)**i
instage_channel_reduce_pow = max_reduction**(
1 / (max(repeats[i] - 1, 1)))
instage_expand_ratio = [
stage_ratio * instage_channel_reduce_pow**x
for x in range(repeats[i])
]
instage_expand_ratios.append(instage_expand_ratio)
print(stage_expand_ratios)
print(instage_expand_ratios)
self._res_blocks = []
channels = [stem_chann] + channels
for idx in range(len(repeats)):
block = X3DTransform(channels[idx],
channels[idx + 1],
kernel_size=3,
stride=strides[idx],
expand_ratio=instage_expand_ratios[idx][0])
self._res_blocks.append(block)
for i in range(repeats[idx] - 1):
block = X3DTransform(
channels[idx + 1],
channels[idx + 1],
kernel_size=3,
stride=1,
se=se[idx],
expand_ratio=instage_expand_ratios[idx][i + 1])
self._res_blocks.append(block)
self._res_blocks = nn.Sequential(*self._res_blocks)
expand_chann = round(channels[-1] * expand_ratio)
self._expand_conv = self.conv3d(channels[-1],
expand_chann,
kernel_size=1,
bias=False)
self._epxand_bn = nn.BatchNorm3d(num_features=expand_chann,
momentum=self._bn_mom,
eps=self._bn_eps)
head_chann = 2048
self._avg_pooling = nn.AdaptiveAvgPool3d(1)
self._dropout = nn.Dropout(0.2)
self._fc1 = nn.Linear(expand_chann, head_chann)
self._fc2 = nn.Linear(head_chann, num_classes)
self._activate = get_act('ReLU')
def __init__(self,
inplanes,
planes,
spatial_stride=1,
temporal_stride=1,
dilation=1,
downsample=None,
style='pytorch',
if_inflate=True,
with_cp=False,
with_trajectory=False):
"""Bottleneck block for ResNet.
If style is "pytorch", the stride-two layer is the 3x3 conv layer,
if it is "caffe", the stride-two layer is the first 1x1 conv layer.
"""
super(Bottleneck, self).__init__()
assert style in ['pytorch', 'caffe']
self.inplanes = inplanes
self.planes = planes
if style == 'pytorch':
self.conv1_stride = 1
self.conv2_stride = spatial_stride
self.conv1_stride_t = 1
self.conv2_stride_t = temporal_stride
else:
self.conv1_stride = spatial_stride
self.conv2_stride = 1
self.conv1_stride_t = temporal_stride
self.conv2_stride_t = 1
self.conv1 = nn.Conv3d(
inplanes,
planes,
kernel_size=1,
stride=(self.conv1_stride_t, self.conv1_stride, self.conv1_stride),
bias=False)
self.conv2 = nn.Conv3d(
planes,
planes,
kernel_size=(1,3,3),
stride=(1, self.conv2_stride, self.conv2_stride),
padding=(0, dilation, dilation),
dilation=(1, dilation, dilation),
bias=False)
self.if_inflate = if_inflate
if self.if_inflate:
self.conv2_t = nn.Conv3d(
planes,
planes,
kernel_size=(3,1,1),
stride=(self.conv2_stride_t,1,1),
padding=(1,0,0),
dilation=1,
bias=True)
self.bn2_t = nn.BatchNorm3d(planes)
self.bn1 = nn.BatchNorm3d(planes)
self.bn2 = nn.BatchNorm3d(planes)
self.conv3 = nn.Conv3d(
planes, planes * self.expansion, kernel_size=1, bias=False)
self.bn3 = nn.BatchNorm3d(planes * self.expansion)
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
self.spatial_tride = spatial_stride
self.temporal_tride = temporal_stride
self.dilation = dilation
self.with_cp = with_cp
self.with_trajectory = with_trajectory
def __init__(self):
super(CNN, self).__init__()
num_filters_layer1 = 32
num_filters_layer2 = 64
num_filters_layer3 = 128
num_filters_layer4 = 256
num_filters_layer5 = 512
kernel_size_layer1 = (7, 11, 11)
kernel_size_layer2 = (5, 9, 9)
kernel_size_layer3 = (5, 7, 7)
kernel_size_layer4 = (3, 5, 5)
kernel_size_layer5 = (3, 3, 3)
filter_stride = 1
filter_padding = 1
pool_size = (2, 3, 3)
pool_stride = 2
pool_padding = 1
self.lrelu = nn.LeakyReLU(inplace=True)
self.maxpool = nn.MaxPool3d(kernel_size=pool_size,
stride=pool_stride,
padding=pool_padding)
# Layer 1
self.conv1 = nn.Conv3d(1,
num_filters_layer1,
kernel_size=kernel_size_layer1,
stride=filter_stride,
padding=filter_padding,
bias=True)
self.bn1 = nn.BatchNorm3d(num_filters_layer1)
# Layer 2
self.conv2 = nn.Conv3d(num_filters_layer1,
num_filters_layer2,
kernel_size=kernel_size_layer2,
stride=filter_stride,
padding=filter_padding,
bias=True)
self.bn2 = nn.BatchNorm3d(num_filters_layer2)
# Layer 3
self.conv3 = nn.Conv3d(num_filters_layer2,
num_filters_layer3,
kernel_size=kernel_size_layer3,
stride=filter_stride,
padding=filter_padding,
bias=True)
self.bn3 = nn.BatchNorm3d(num_filters_layer3)
# Layer 4
self.conv4 = nn.Conv3d(num_filters_layer3,
num_filters_layer4,
kernel_size=kernel_size_layer4,
stride=filter_stride,
padding=filter_padding,
bias=True)
self.bn4 = nn.BatchNorm3d(num_filters_layer4)
# Layer 5
self.conv5 = nn.Conv3d(num_filters_layer4,
num_filters_layer5,
kernel_size=kernel_size_layer5,
stride=filter_stride,
padding=filter_padding,
bias=True)
self.bn5 = nn.BatchNorm3d(num_filters_layer5)
# Global Average Pooling Layer
self.aap = nn.AdaptiveAvgPool3d(output_size=(1, 1, 1))
# Linear layer
self.fc = nn.Linear(num_filters_layer5, 1)
def __init__(self, column_units):
super(Model, self).__init__()
self.block1 = nn.Sequential(
nn.Conv3d(3,
32,
kernel_size=(3, 5, 5),
stride=(1, 2, 2),
dilation=(1, 1, 1),
padding=(1, 2, 2)),
nn.BatchNorm3d(32),
nn.ReLU(inplace=True),
nn.Dropout3d(p=0.2),
)
self.block2 = nn.Sequential(
nn.Conv3d(32,
64,
kernel_size=(3, 3, 3),
stride=1,
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(64),
nn.ReLU(inplace=True),
nn.Conv3d(64,
128,
kernel_size=(3, 3, 3),
stride=(1, 2, 2),
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(128),
nn.ReLU(inplace=True),
nn.Dropout3d(p=0.2),
)
self.block3 = nn.Sequential(
nn.Conv3d(128,
128,
kernel_size=(3, 3, 3),
stride=1,
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(128),
nn.ReLU(inplace=True),
nn.Conv3d(128,
128,
kernel_size=(3, 3, 3),
stride=1,
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(128),
nn.ReLU(inplace=True),
nn.Conv3d(128,
256,
kernel_size=(3, 3, 3),
stride=(1, 2, 2),
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(256),
nn.ReLU(inplace=True),
nn.Dropout3d(p=0.2),
)
self.block4 = nn.Sequential(
nn.Conv3d(256,
256,
kernel_size=(3, 3, 3),
stride=1,
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(256),
nn.ReLU(inplace=True),
nn.Conv3d(256,
256,
kernel_size=(3, 3, 3),
stride=1,
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(256),
nn.ReLU(inplace=True),
nn.Conv3d(256,
512,
kernel_size=(3, 3, 3),
stride=(1, 2, 2),
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(512),
nn.ReLU(inplace=True),
nn.Dropout3d(p=0.2),
)
self.block5 = nn.Sequential(
nn.Conv3d(512,
512,
kernel_size=(3, 3, 3),
stride=1,
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(512),
nn.ReLU(inplace=True),
nn.Conv3d(512,
512,
kernel_size=(3, 3, 3),
stride=(1, 2, 2),
dilation=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(512),
nn.ReLU(inplace=True),
)
def __init__(self, in_channels, out_channels, k_size=3, stride=1, padding=1):
super(ConvBlock, self).__init__()
self.conv3d = nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=k_size,
stride=stride, padding=padding)
self.batch_norm = nn.BatchNorm3d(num_features=out_channels)
def __init__(self,
A=32,
B=32,
C=32,
D=32,
E=10,
K=3,
P=4,
iters=3,
device='cuda',
_lambda=[]):
super(MatrixCapsules, self).__init__()
self.A, self.B, self.C, self.D, self.E, self.P = A, B, C, D, E, P
self.conv1 = nn.Conv2d(in_channels=1,
out_channels=A,
kernel_size=5,
stride=2,
padding=2)
self.relu1 = nn.ReLU(inplace=False)
self.primary_caps = PrimaryCaps(A, B, 1, P, stride=1)
self.conv_caps1 = ConvCaps(B,
C,
K,
P,
stride=2,
iters=iters,
device=device,
_lambda=_lambda[0])
self.conv_caps2 = ConvCaps(C,
D,
K,
P,
stride=1,
iters=iters,
device=device,
_lambda=_lambda[1])
self.class_caps = ConvCaps(D,
E,
1,
P,
stride=1,
iters=iters,
device=device,
_lambda=_lambda[2],
coor_add=True,
w_shared=True)
self.batch_norm_input = nn.BatchNorm2d(num_features=A, affine=True)
self.drop_out_input = nn.Dropout2d(p=.2)
self.batch_norm_3d_1 = nn.BatchNorm3d(B, affine=True)
self.batch_norm_2d_1 = nn.BatchNorm2d(B, affine=True)
self.batch_norm_3d_2 = nn.BatchNorm3d(C, affine=True)
self.batch_norm_2d_2 = nn.BatchNorm2d(C, affine=True)
self.batch_norm_3d_3 = nn.BatchNorm3d(D, affine=True)
self.batch_norm_2d_3 = nn.BatchNorm2d(D, affine=True)
self.drop_out_3d_1 = nn.Dropout(p=.1)
self.drop_out_2d_1 = nn.Dropout(p=.1)
self.drop_out_3d_2 = nn.Dropout(p=.1)
self.drop_out_2d_2 = nn.Dropout(p=.1)
self.drop_out_3d_3 = nn.Dropout(p=.1)
self.drop_out_2d_3 = nn.Dropout(p=.1)
self.to(device)
def __init__(self, num_channels, running_mean, running_var):
super(SimpleBatchNorm, self).__init__()
self.batchnorm = nn.BatchNorm3d(num_channels)
self.batchnorm.running_mean = running_mean
self.batchnorm.running_var = running_var
def __init__(self,
in_channels,
out_channels,
bn=True,
bias=False,
sample="none-3",
act_fun="ReLU",
drop=0.):
super(Partial3DConv, self).__init__()
if sample == "down-7":
# Kernel Size = 7, Stride = 2, Padding = 3
self.input_conv = nn.Conv3d(in_channels,
out_channels,
7,
2,
3,
bias=bias)
self.mask_conv = nn.Conv3d(in_channels,
out_channels,
7,
2,
3,
bias=False)
elif sample == "down-5":
self.input_conv = nn.Conv3d(in_channels,
out_channels,
5,
2,
2,
bias=bias)
self.mask_conv = nn.Conv3d(in_channels,
out_channels,
5,
2,
2,
bias=False)
elif sample == "down-3":
self.input_conv = nn.Conv3d(in_channels,
out_channels,
3,
2,
1,
bias=bias)
self.mask_conv = nn.Conv3d(in_channels,
out_channels,
3,
2,
1,
bias=False)
else:
self.input_conv = nn.Conv3d(in_channels,
out_channels,
3,
1,
1,
bias=bias)
self.mask_conv = nn.Conv3d(in_channels,
out_channels,
3,
1,
1,
bias=False)
nn.init.constant_(self.mask_conv.weight, 1.0)
# "Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification"
# negative slope of leaky_relu set to 0, same as relu
# "fan_in" preserved variance from forward pass
nn.init.kaiming_normal_(self.input_conv.weight, a=0, mode="fan_in")
for param in self.mask_conv.parameters():
param.requires_grad = False
self.bn = nn.BatchNorm3d(out_channels) if bn else None
self.act = act(act_fun)
self.dr = nn.Dropout2d(drop)
def test_resnet3d_bottleneck():
data = torch.randn(1, 256, 1, 56, 56)
inplanes = 256
planes = 128
expansion = ResNet3DBottleneck.expansion
# 不膨胀
# 不进行下采样
temporal_stride = 1
spatial_stride = 1
downsample = nn.Sequential(
nn.Conv3d(inplanes,
planes * expansion,
kernel_size=1,
stride=(temporal_stride, spatial_stride, spatial_stride),
bias=False),
nn.BatchNorm3d(planes * expansion),
)
model = ResNet3DBottleneck(in_planes=inplanes,
out_planes=planes,
spatial_stride=spatial_stride,
temporal_stride=temporal_stride,
inflate=False,
down_sample=downsample)
print(model)
outputs = model(data)
print(outputs.shape)
assert outputs.shape == (1, planes * expansion, 1, 56, 56)
# 进行下采样
temporal_stride = 1
spatial_stride = 2
downsample = nn.Sequential(
nn.Conv3d(inplanes,
planes * expansion,
kernel_size=1,
stride=(temporal_stride, spatial_stride, spatial_stride),
bias=False),
nn.BatchNorm3d(planes * expansion),
)
model = ResNet3DBottleneck(in_planes=inplanes,
out_planes=planes,
spatial_stride=spatial_stride,
temporal_stride=temporal_stride,
inflate=False,
down_sample=downsample)
print(model)
outputs = model(data)
print(outputs.shape)
assert outputs.shape == (1, planes * expansion, 1, 28, 28)
# 32x4d
# 进行下采样
temporal_stride = 1
spatial_stride = 2
downsample = nn.Sequential(
nn.Conv3d(inplanes,
planes * expansion,
kernel_size=1,
stride=(temporal_stride, spatial_stride, spatial_stride),
bias=False),
nn.BatchNorm3d(planes * expansion),
)
model = ResNet3DBottleneck(in_planes=inplanes,
out_planes=planes,
spatial_stride=spatial_stride,
temporal_stride=temporal_stride,
inflate=False,
down_sample=downsample,
groups=32,
base_width=4)
print(model)
outputs = model(data)
print(outputs.shape)
assert outputs.shape == (1, planes * expansion, 1, 28, 28)
print("iter:{}, bias.shape: {}".format(i, bn.bias.shape))
# %% nn.BatchNorm3d
batch_size = 3
num_features = 4
momentum = 0.3
features_shape = (2, 2, 3)
feature = torch.ones(features_shape) # 3D
feature_map = torch.stack([feature * (i + 1) for i in range(num_features)],
dim=0) # 4D
feature_maps = torch.stack([feature_map for i in range(batch_size)],
dim=0) # 5D
print("input data:\n{} shape is {}".format(feature_maps, feature_maps.shape))
bn = nn.BatchNorm3d(num_features=num_features, momentum=momentum)
running_mean, running_var = 0, 1
for i in range(2):
outputs = bn(feature_maps)
print("\niter:{}, running_mean.shape: {}".format(i, bn.running_mean.shape))
print("iter:{}, running_var.shape: {}".format(i, bn.running_var.shape))
print("iter:{}, weight.shape: {}".format(i, bn.weight.shape))
print("iter:{}, bias.shape: {}".format(i, bn.bias.shape))
def __init__(self, block, layers, num_classes=1):
self.inplanes = 64
self.atrous_depth = 256
super(ResNet, self).__init__()
self.conv1 = nn.Conv3d(1,
64,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.bn1 = nn.BatchNorm3d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool3d(kernel_size=3, stride=2,
padding=1) # change
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
self.layer3 = self._make_layer(block,
256,
layers[2],
stride=1,
dilation=1) # deeplab change
self.layer4 = self._make_layer(block,
512,
layers[3],
stride=1,
dilation=1) # deeplab change
self.fc_aspp0 = nn.Conv3d(512 * block.expansion,
self.atrous_depth,
kernel_size=1,
stride=1,
padding=0) # deelab v3
self.fc_aspp1 = nn.Conv3d(512 * block.expansion,
self.atrous_depth,
kernel_size=3,
stride=1,
dilation=6,
padding=6) # deeplab change
self.fc_aspp2 = nn.Conv3d(512 * block.expansion,
self.atrous_depth,
kernel_size=3,
stride=1,
dilation=12,
padding=12) # deeplab change
self.fc_aspp3 = nn.Conv3d(512 * block.expansion,
self.atrous_depth,
kernel_size=3,
stride=1,
dilation=18,
padding=18) # deeplab change
self.fc_aspp_cat = nn.Conv3d(self.atrous_depth * 4,
self.atrous_depth,
kernel_size=1,
stride=1,
padding=0)
self.fc_dropout = nn.Dropout3d(0.9, inplace=True)
self.fc_bn = nn.BatchNorm3d(self.atrous_depth)
self.fc_aspp_output = nn.Conv3d(self.atrous_depth,
num_classes,
kernel_size=1,
stride=1,
padding=0)
#self.fc_aspp_cat = nn.Conv3d(self.atrous_depth, num_classes, kernel_size=1, stride=1, padding=0)
for m in self.modules():
if isinstance(m, nn.Conv3d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm3d):
m.weight.data.fill_(1)
m.bias.data.zero_()
def __init__(self, with_classifier=True, num_classes=101):
super(C3D, self).__init__()
self.with_classifier = with_classifier
self.num_classes = num_classes
self.conv1 = nn.Conv3d(3, 64, kernel_size=(3, 3, 3), padding=(1, 1, 1))
self.bn1 = nn.BatchNorm3d(64)
self.relu1 = nn.ReLU()
self.pool1 = nn.MaxPool3d(kernel_size=(1, 2, 2), stride=(1, 2, 2))
self.conv2 = nn.Conv3d(64,
128,
kernel_size=(3, 3, 3),
padding=(1, 1, 1))
self.bn2 = nn.BatchNorm3d(128)
self.relu2 = nn.ReLU()
self.pool2 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2))
self.conv3a = nn.Conv3d(128,
256,
kernel_size=(3, 3, 3),
padding=(1, 1, 1))
self.bn3a = nn.BatchNorm3d(256)
self.relu3a = nn.ReLU()
self.conv3b = nn.Conv3d(256,
256,
kernel_size=(3, 3, 3),
padding=(1, 1, 1))
self.bn3b = nn.BatchNorm3d(256)
self.relu3b = nn.ReLU()
self.pool3 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2))
self.conv4a = nn.Conv3d(256,
512,
kernel_size=(3, 3, 3),
padding=(1, 1, 1))
self.bn4a = nn.BatchNorm3d(512)
self.relu4a = nn.ReLU()
self.conv4b = nn.Conv3d(512,
512,
kernel_size=(3, 3, 3),
padding=(1, 1, 1))
self.bn4b = nn.BatchNorm3d(512)
self.relu4b = nn.ReLU()
self.pool4 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=(2, 2, 2))
self.conv5a = nn.Conv3d(512,
512,
kernel_size=(3, 3, 3),
padding=(1, 1, 1))
self.bn5a = nn.BatchNorm3d(512)
self.relu5a = nn.ReLU()
self.conv5b = nn.Conv3d(512,
512,
kernel_size=(3, 3, 3),
padding=(1, 1, 1))
self.bn5b = nn.BatchNorm3d(512)
self.relu5b = nn.ReLU()
self.pool5 = nn.AdaptiveAvgPool3d(1)
if self.with_classifier:
self.linear = nn.Linear(512, self.num_classes)
def __init__(self, in_planes, out_planes):
super(Deconv3DBlock, self).__init__()
self.block = nn.Sequential(
SingleDeconv3DBlock(in_planes, out_planes),
SingleConv3DBlock(out_planes, out_planes, 3),
nn.BatchNorm3d(out_planes), nn.ReLU(True))
def __init__(self, in_planes, out_planes, kernel_size):
super(Conv3DBlock, self).__init__()
self.block = nn.Sequential(
SingleConv3DBlock(in_planes, out_planes, kernel_size),
nn.BatchNorm3d(out_planes), nn.ReLU(True))
def __init__(self, args):
super(PSMNet, self).__init__()
self.maxdisp = args.maxdisp
self.planes = args.planes
inplanes = self.planes * 2
self.feature_extraction = feature_extraction(args)
if args.shuffle:
block3d = ResBlock3DShuffle
else:
block3d = ResBlock3D
self.fuse = nn.Sequential(
nn.Conv3d(inplanes,
self.planes,
kernel_size=3,
stride=1,
padding=1,
dilation=1,
bias=False),
nn.BatchNorm3d(self.planes),
nn.ReLU(inplace=True),
)
self.unet_conv = nn.ModuleList()
inplanes = self.planes
for i in range(3):
outplanes = inplanes * 2
self.unet_conv.append(
nn.Sequential(
nn.Conv3d(inplanes,
outplanes,
kernel_size=3,
stride=2,
padding=1,
dilation=1,
bias=False),
nn.BatchNorm3d(outplanes),
nn.ReLU(inplace=True),
))
inplanes = outplanes
self.classifiers = nn.ModuleList()
for i in range(4):
outplanes = inplanes // 2
self.classifiers.append(
nn.Sequential(
nn.ConvTranspose3d(inplanes,
1,
kernel_size=7,
stride=4,
padding=3,
output_padding=3,
dilation=1,
bias=False)))
inplanes = outplanes
self.disparityregression = disparityregression(self.maxdisp)
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.Conv3d):
n = m.kernel_size[0] * m.kernel_size[1] * m.kernel_size[
2] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.BatchNorm3d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.Linear):
m.bias.data.zero_()
def test_resnet3d_bottleneck_3_1_1():
data = torch.randn(1, 64, 8, 56, 56)
inplanes = 64
planes = 128
expansion = ResNet3DBottleneck.expansion
inflate_style = '3x1x1'
# 膨胀，不进行时间下采样
# 不进行下采样
temporal_stride = 1
spatial_stride = 1
downsample = nn.Sequential(
nn.Conv3d(inplanes,
planes * expansion,
kernel_size=1,
stride=(temporal_stride, spatial_stride, spatial_stride),
bias=False),
nn.BatchNorm3d(planes * expansion),
)
model = ResNet3DBottleneck(in_planes=inplanes,
out_planes=planes,
spatial_stride=spatial_stride,
temporal_stride=temporal_stride,
inflate=True,
inflate_style=inflate_style,
down_sample=downsample)
print(model)
outputs = model(data)
print(outputs.shape)
assert outputs.shape == (1, planes * expansion, 8, 56, 56)
# 膨胀，进行时间下采样
# 不进行下采样
temporal_stride = 2
spatial_stride = 1
downsample = nn.Sequential(
nn.Conv3d(inplanes,
planes * expansion,
kernel_size=1,
stride=(temporal_stride, spatial_stride, spatial_stride),
bias=False),
nn.BatchNorm3d(planes * expansion),
)
model = ResNet3DBottleneck(in_planes=inplanes,
out_planes=planes,
spatial_stride=spatial_stride,
temporal_stride=temporal_stride,
inflate=True,
inflate_style=inflate_style,
down_sample=downsample)
print(model)
outputs = model(data)
print(outputs.shape)
assert outputs.shape == (1, planes * expansion, 4, 56, 56)
# 膨胀，进行时间下采样
# 进行下采样
temporal_stride = 2
spatial_stride = 2
downsample = nn.Sequential(
nn.Conv3d(inplanes,
planes * expansion,
kernel_size=1,
stride=(temporal_stride, spatial_stride, spatial_stride),
bias=False),
nn.BatchNorm3d(planes * expansion),
)
model = ResNet3DBottleneck(in_planes=inplanes,
out_planes=planes,
spatial_stride=spatial_stride,
temporal_stride=temporal_stride,
inflate=True,
inflate_style=inflate_style,
down_sample=downsample)
print(model)
outputs = model(data)
print(outputs.shape)
assert outputs.shape == (1, planes * expansion, 4, 28, 28)
def __init__(self,
in_channels,
out_channels,
kernel_size,
with_se=False,
normalize=True,
num_cls=3,
num_scale=5):
super().__init__()
self.num_scale = num_scale
self.convblocks_1 = ConvBlock(in_channels, out_channels,
kernel_size).cuda()
self.convblocks_2 = ConvBlock(in_channels, out_channels,
kernel_size).cuda()
self.convblocks_3 = ConvBlock(in_channels, out_channels,
kernel_size).cuda()
self.convblocks_4 = ConvBlock(in_channels, out_channels,
kernel_size).cuda()
self.convblocks_5 = ConvBlock(in_channels, out_channels,
kernel_size).cuda()
# for i in range(num_scale):
# self.convblocks.append(ConvBlock(in_channels, out_channels, kernel_size).cuda())
# voxel_layers = [
# nn.Conv3d(in_channels, out_channels, kernel_size, stride=1, padding=kernel_size // 2),
# nn.BatchNorm3d(out_channels, eps=1e-4),
# nn.LeakyReLU(0.1, True),
# nn.MaxPool3d(3),
# nn.Conv3d(out_channels, out_channels, kernel_size, stride=1, padding=kernel_size // 2),
# nn.BatchNorm3d(out_channels, eps=1e-4),
# nn.LeakyReLU(0.1, True),
# nn.MaxPool3d(3),
# nn.Conv3d(out_channels, out_channels, kernel_size, stride=1, padding=kernel_size // 2),
# nn.BatchNorm3d(out_channels, eps=1e-4),
# nn.LeakyReLU(0.1, True),
# nn.MaxPool3d(3),
# ]
# self.voxel_layers = []
# self.voxel_layers.append(nn.Sequential(*voxel_layers).cuda())
# if (self.num_scale > 2):
# for i in range(self.num_scale - 1):
# print(i)
# self.voxel_layers.append(nn.Sequential(*voxel_layers).cuda())
concate_layers = [
nn.Conv3d(out_channels * num_scale, out_channels, 1, stride=1),
nn.BatchNorm3d(out_channels, eps=1e-4),
nn.LeakyReLU(0.1, True),
]
self.concate_layers = nn.Sequential(*concate_layers)
# self.fc_lyaer = nn.Sequential(
# nn.Conv1d(64, 64, kernel_size=1, bias=False),
# nn.BatchNorm1d(64),
# nn.ReLU(True),
# nn.Dropout(0.5),
# nn.Conv1d(64, 2, kernel_size=1),
# )
self.fc = nn.Sequential(
nn.Linear(32, 32, bias=False),
nn.BatchNorm1d(32),
nn.ReLU(True),
nn.Dropout(0.5),
nn.Linear(32, num_cls, bias=False),
nn.Softmax(1),
)
def __init__(self, nlabel, mindepth):
super(YunNetTest, self).__init__()
self.nlabel = nlabel
self.mindepth = mindepth
self.cdcor = CDcor(64, 32)
#spp
self.feature_extraction = feature_extraction()
#3DCNN
self.conv3d0 = nn.Sequential(
convbn_3d(64, 32, 3, 1, 1),
nn.ReLU(inplace=True),
convbn_3d(32, 32, 3, 1, 1),
nn.ReLU(inplace=True),
)
self.conv3d0_2 = nn.Sequential(
convbn_3d(32, 32, 1, 1, 0),
nn.ReLU(inplace=True),
)
self.conv3d1 = nn.Sequential(
convbn_3d(32, 64, 3, 2, 1),
nn.ReLU(inplace=True),
convbn_3d(64, 64, 3, 1, 1),
)
self.conv3d1_2 = nn.Sequential(
convbn_3d(64, 64, 1, 1, 0),
nn.ReLU(inplace=True),
)
self.conv3d2 = nn.Sequential(
convbn_3d(64, 128, 3, 2, 1),
nn.ReLU(inplace=True),
convbn_3d(128, 128, 3, 1, 1),
)
self.conv3d2_2 = nn.Sequential(
convbn_3d(128, 128, 1, 1, 0),
nn.ReLU(inplace=True),
)
self.conv2_3 = nn.Sequential(
nn.ConvTranspose3d(128,
64,
kernel_size=3,
padding=1,
output_padding=1,
stride=2,
bias=False), nn.BatchNorm3d(64))
self.conv3_2 = nn.Sequential(
nn.ConvTranspose3d(64,
32,
kernel_size=3,
padding=1,
output_padding=1,
stride=2,
bias=False), nn.BatchNorm3d(32))
self.cost_disp1 = nn.Sequential(
nn.Upsample(scale_factor=2, mode='trilinear'),
nn.ReLU(inplace=True),
convbn_3d(64, 1, 3, 1, 1),
)
self.cost_disp2 = nn.Sequential(
convbn_3d(32, 32, 3, 1, 1),
nn.ReLU(inplace=True),
convbn_3d(32, 1, 3, 1, 1),
)
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.Conv3d):
n = m.kernel_size[0] * m.kernel_size[1] * m.kernel_size[
2] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.BatchNorm3d):
m.weight.data.fill_(1)
m.bias.data.zero_()
def _make_conv_layer(self, in_c, out_c, pool_size, stride):
conv_layer = nn.Sequential(
nn.Conv3d(in_c, out_c, kernel_size=3, stride=1, padding=1),
nn.BatchNorm3d(out_c), nn.ELU(),
nn.MaxPool3d(pool_size, stride=stride, padding=0))
return conv_layer
def __init__(self, cube_len=64):
super(UNet, self).__init__()
self.cube_len = cube_len
self.code_len = cube_len * 8
#Contracting path:
self.enc_1 = nn.Sequential(
nn.Conv3d(1, self.cube_len, kernel_size = 4, stride = 2, padding = 1),
nn.BatchNorm3d(self.cube_len),
nn.ReLU()
)
self.enc_2 = nn.Sequential(
nn.Conv3d(self.cube_len, self.cube_len * 2, kernel_size = 4, stride = 2, padding = 1),
nn.BatchNorm3d(self.cube_len * 2),
nn.ReLU()
)
self.enc_3 = nn.Sequential(
nn.Conv3d(self.cube_len * 2, self.cube_len * 4, kernel_size = 4, stride = 2, padding = 1),
nn.BatchNorm3d(self.cube_len * 4),
nn.ReLU()
)
self.enc_4 = nn.Sequential(
nn.Conv3d(self.cube_len * 4, self.code_len, kernel_size = 4, stride = 2, padding = 1),
nn.BatchNorm3d(self.code_len),
nn.ReLU()
)
self.enc_5 = nn.Sequential(
nn.Conv3d(self.code_len, self.code_len, kernel_size = 4, stride = 2, padding = 1),
nn.BatchNorm3d(self.code_len),
nn.ReLU()
)
self.enc_6 = nn.Sequential(
nn.Conv3d(self.code_len, self.code_len, kernel_size = 4, stride = 2, padding = 1),
#cant batch norm when features are 1x1x1
#nn.BatchNorm3d(self.code_len),
nn.ReLU()
)
#Expansive path
self.dec_1 = torch.nn.Sequential(
nn.ConvTranspose3d(self.code_len, self.code_len, kernel_size=4, stride=2, padding = 1),
nn.BatchNorm3d(self.code_len),
nn.ReLU()
)
#According to the paper this layer also has dropout
self.dec_2 = torch.nn.Sequential(
nn.ConvTranspose3d( (self.code_len) * 2, self.code_len, kernel_size=4, stride=2, padding = 1),
nn.BatchNorm3d(self.code_len),
nn.ReLU()
)
self.dec_3 = torch.nn.Sequential(
nn.ConvTranspose3d( (self.code_len) * 2, (self.cube_len * 4) , kernel_size=4, stride=2, padding = 1),
nn.BatchNorm3d((self.cube_len * 4)),
nn.ReLU()
)
self.dec_4 = torch.nn.Sequential(
nn.ConvTranspose3d( (self.cube_len * 4) * 2, (self.cube_len * 2) , kernel_size=4, stride=2, padding = 1),
nn.BatchNorm3d((self.cube_len * 2)),
nn.ReLU()
)
self.dec_5 = torch.nn.Sequential(
nn.ConvTranspose3d( (self.cube_len * 2) * 2, self.cube_len, kernel_size=4, stride=2, padding = 1),
nn.BatchNorm3d(self.cube_len),
nn.ReLU()
)
self.dec_6 = torch.nn.Sequential(
nn.ConvTranspose3d( (self.cube_len) * 2, 1, kernel_size=4, stride=2, padding = 1),
nn.BatchNorm3d(1),
nn.ReLU()
)
def __init__(self):
super(BasicStem, self).__init__(
nn.Conv3d(3, 64, kernel_size=(3, 7, 7), stride=(1, 2, 2),
padding=(1, 3, 3), bias=False),
nn.BatchNorm3d(64),
nn.ReLU(inplace=True))
def __init__(self,
block=Bottleneck,
layers=[3, 4, 6, 3],
class_num=10,
dropout=0.5):
super(SlowFast, self).__init__()
self.fast_inplanes = 8
self.fast_conv1 = nn.Conv3d(3,
8,
kernel_size=(5, 7, 7),
stride=(1, 2, 2),
padding=(2, 3, 3),
bias=False)
self.fast_bn1 = nn.BatchNorm3d(8)
self.fast_relu = nn.ReLU(inplace=True)
self.fast_maxpool = nn.MaxPool3d(kernel_size=(1, 3, 3),
stride=(1, 2, 2),
padding=(0, 1, 1))
self.fast_res2 = self._make_layer_fast(block,
8,
layers[0],
head_conv=3)
self.fast_res3 = self._make_layer_fast(block,
16,
layers[1],
stride=2,
head_conv=3)
self.fast_res4 = self._make_layer_fast(block,
32,
layers[2],
stride=2,
head_conv=3)
self.fast_res5 = self._make_layer_fast(block,
64,
layers[3],
stride=2,
head_conv=3)
self.lateral_p1 = nn.Conv3d(8,
8 * 2,
kernel_size=(7, 1, 1),
stride=(8, 1, 1),
bias=False,
padding=(3, 0, 0))
self.lateral_res2 = nn.Conv3d(32,
32 * 2,
kernel_size=(7, 1, 1),
stride=(8, 1, 1),
bias=False,
padding=(3, 0, 0))
self.lateral_res3 = nn.Conv3d(64,
64 * 2,
kernel_size=(7, 1, 1),
stride=(8, 1, 1),
bias=False,
padding=(3, 0, 0))
self.lateral_res4 = nn.Conv3d(128,
128 * 2,
kernel_size=(7, 1, 1),
stride=(8, 1, 1),
bias=False,
padding=(3, 0, 0))
self.lateral_p1_bn = nn.BatchNorm3d(16)
self.lateral_res2_bn = nn.BatchNorm3d(64)
self.lateral_res3_bn = nn.BatchNorm3d(128)
self.lateral_res4_bn = nn.BatchNorm3d(256)
self.slow_inplanes = 64 + 64 // 8 * 2
self.slow_conv1 = nn.Conv3d(3,
64,
kernel_size=(1, 7, 7),
stride=(1, 2, 2),
padding=(0, 3, 3),
bias=False)
self.slow_bn1 = nn.BatchNorm3d(64)
self.slow_relu = nn.ReLU(inplace=True)
self.slow_maxpool = nn.MaxPool3d(kernel_size=(1, 3, 3),
stride=(1, 2, 2),
padding=(0, 1, 1))
self.slow_res2 = self._make_layer_slow(block,
64,
layers[0],
head_conv=1)
self.slow_res3 = self._make_layer_slow(block,
128,
layers[1],
stride=2,
head_conv=1)
self.slow_res4 = self._make_layer_slow(block,
256,
layers[2],
stride=2,
head_conv=3)
self.slow_res5 = self._make_layer_slow(block,
512,
layers[3],
stride=2,
head_conv=3)
def conv_1x1x1_bn(inp, oup):
return nn.Sequential(nn.Conv3d(inp, oup, 1, 1, 0, bias=False),
nn.BatchNorm3d(oup), nn.ReLU6(inplace=True))
def __init__(self, seg2d_path, img_size):
super(ImageGen3DNet, self).__init__()
self.seg2d = Seg2DNet(model=models.resnet101(False), num_classes=12)
chpo = torch.load(seg2d_path)
self.seg2d.load_state_dict(chpo['state_dict'], strict=False)
print "=> seg2d loaded checkpoint '{}'".format(seg2d_path)
self.seq1 = nn.Sequential(nn.Conv3d(64, 64, 3, padding=1, bias=False),
nn.BatchNorm3d(64), nn.ReLU(inplace=True),
nn.Conv3d(64, 64, 3, padding=1, bias=False),
nn.BatchNorm3d(64))
self.seq2 = nn.Sequential(nn.Conv3d(64, 64, 3, padding=1, bias=False),
nn.BatchNorm3d(64), nn.ReLU(inplace=True),
nn.Conv3d(64, 64, 3, padding=1, bias=False),
nn.BatchNorm3d(64))
self.relu = nn.ReLU(inplace=True)
self.ASPP3D1 = ASPP3D(64, 64, [1, 2, 3])
self.ASPP3D2 = ASPP3D(64, 64, [1, 2, 3])
self.ASPP3Dout = nn.Sequential(nn.Conv3d(256, 128, 1, bias=False),
nn.BatchNorm3d(128),
nn.ReLU(inplace=True),
nn.Conv3d(128, 128, 1, bias=False),
nn.BatchNorm3d(128),
nn.ReLU(inplace=True),
nn.Conv3d(128, 12, 1),
nn.Conv3d(12, 12, 3, padding=1))
self.img_required_size = (640, 480)
self.img_size = img_size
if cmp(self.img_required_size, self.img_size) != 0:
x = np.array(range(self.img_required_size[0]), dtype=np.float32)
y = np.array(range(self.img_required_size[1]), dtype=np.float32)
scale = 1.0 * self.img_size[0] / self.img_required_size[0]
x = x * scale + 0.5
y = y * scale + 0.5
x = x.astype(np.int64)
y = y.astype(np.int64)
if x[self.img_required_size[0] - 1] >= self.img_size[0]:
x[self.img_required_size[0] - 1] = self.img_size[0] - 1
if y[self.img_required_size[1] - 1] >= self.img_size[1]:
y[self.img_required_size[1] - 1] = self.img_size[1] - 1
xx = np.ones(
(self.img_required_size[1], self.img_required_size[0]),
dtype=np.int64)
yy = np.ones(
(self.img_required_size[1], self.img_required_size[0]),
dtype=np.int64)
xx[:] = x
yy[:] = y.reshape(
(self.img_required_size[1], 1)) * self.img_size[0]
self.image_mapping = (xx + yy).reshape(-1)
else:
self.image_mapping = np.array(range(self.img_required_size[0] *
self.img_required_size[1]),
dtype=np.int64)
self.image_mapping = torch.autograd.Variable(torch.cuda.LongTensor(
self.image_mapping),
requires_grad=False)
self.dim_inc_dim = 64
def __init__(self,
block,
layers,
sample_input_D,
sample_input_H,
sample_input_W,
num_classes,
shortcut_type='B',
no_cuda=False):
self.inplanes = 64
self.no_cuda = no_cuda
super(ResNet, self).__init__()
self.conv1 = nn.Conv3d(1,
64,
kernel_size=7,
stride=(2, 2, 2),
padding=(3, 3, 3),
bias=False)
self.bn1 = nn.BatchNorm3d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool3d(kernel_size=(3, 3, 3), stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0], shortcut_type)
self.layer2 = self._make_layer(block,
128,
layers[1],
shortcut_type,
stride=2)
self.layer3 = self._make_layer(block,
256,
layers[2],
shortcut_type,
stride=1,
dilation=2)
self.layer4 = self._make_layer(block,
512,
layers[3],
shortcut_type,
stride=1,
dilation=4)
# segmentation module is voided
# self.conv_seg = nn.Sequential( nn.ConvTranspose3d(512 * block.expansion,32,2,stride=2),
# nn.BatchNorm3d(32),
# nn.ReLU(inplace=True),
# nn.Conv3d(32,32,kernel_size=3,stride=(1, 1, 1),padding=(1, 1, 1),bias=False),
# nn.BatchNorm3d(32),
# nn.ReLU(inplace=True),
# nn.Conv3d(32,num_classes,kernel_size=1,stride=(1, 1, 1),bias=False)
# )
# FCN Classifier Module for IQA
self.classifier = nn.Sequential(nn.Linear(512 * 8 * 8 * 8, 4096),
nn.ReLU(inplace=True), nn.Dropout(0.5),
nn.Linear(4096, 1024),
nn.ReLU(inplace=True),
nn.Linear(1024, num_classes))
for m in self.modules():
if isinstance(m, nn.Conv3d):
m.weight = nn.init.kaiming_normal_(m.weight, mode='fan_out')
elif isinstance(m, nn.BatchNorm3d):
m.weight.data.fill_(1)
m.bias.data.zero_()
def __init__(self, block, layers, opt):
# default
shortcut_type = 'B'
cardinality = 32
num_classes = 400
# user paras
num_classes = opt.n_classes
shortcut_type = opt.resnet_shortcut
cardinality = opt.resnext_cardinality
sample_size = opt.sample_size
sample_duration = opt.basic_duration
self.learning_policy = opt.learning_policy
self.num_classes = opt.n_classes
self.inplanes = 64
super(ResNeXt, self).__init__()
down_stride_1 = (1, 2, 2)
down_stride_2 = (2, 2, 2)
self.conv1 = nn.Conv3d(3,
64,
kernel_size=7,
stride=down_stride_1,
padding=(3, 3, 3),
bias=False)
self.bn1 = nn.BatchNorm3d(64)
self.relu = nn.ReLU(inplace=True)
base_c = 128
self.maxpool = nn.MaxPool3d(kernel_size=(3, 3, 3),
stride=down_stride_2,
padding=1)
self.layer1 = self._make_layer(block, base_c, layers[0], shortcut_type,
cardinality)
self.layer2 = self._make_layer(block,
base_c * 2,
layers[1],
shortcut_type,
cardinality,
stride=down_stride_2)
self.layer3 = self._make_layer(block,
base_c * 4,
layers[2],
shortcut_type,
cardinality,
stride=down_stride_2)
self.layer4 = self._make_layer(block,
base_c * 8,
layers[3],
shortcut_type,
cardinality,
stride=down_stride_2)
last_duration = int(1)
last_size = int(math.ceil(sample_size / 32.0))
self.avgpool = nn.AvgPool3d((last_duration, last_size, last_size),
stride=1)
self.t_all = int(sample_duration / 16.0)
self.dims = int(base_c * 8 * block.expansion)
self.fc_emd = self.dims * self.t_all / 2
if self.learning_policy == '2stream':
self.fc_cls_1 = nn.Linear(self.fc_emd, 2 * num_classes).cuda()
self.fc_box_1 = nn.Linear(self.fc_emd, num_classes).cuda()
self.fc_cls_2 = nn.Linear(self.fc_emd, 2 * num_classes).cuda()
self.fc_box_2 = nn.Linear(self.fc_emd, num_classes).cuda()
self.fc = []
self.others = []
for m in self.modules():
if isinstance(m, nn.Conv3d):
m.weight = nn.init.kaiming_normal(m.weight, mode='fan_out')
elif isinstance(m, nn.BatchNorm3d):
m.weight.data.fill_(1)
m.bias.data.zero_()
def __init__(self,
depth,
pretrained=None,
num_stages=4,
spatial_strides=(1, 2, 2, 2),
temporal_strides=(1, 1, 1, 1),
dilations=(1, 1, 1, 1),
out_indices=(0, 1, 2, 3),
conv1_kernel_t=5,
conv1_stride_t=2,
pool1_kernel_t=1,
pool1_stride_t=2,
use_pool2=True,
style='pytorch',
frozen_stages=-1,
inflate_freq=(1, 1, 1, 1), # For C2D baseline, this is set to -1.
bn_eval=True,
bn_frozen=False,
partial_bn=False,
with_cp=False,
with_trajectory=False,
trajectory_source_indices=-1,
trajectory_downsample_method='ave',
conv_bias=0.2):
super(ResNet_S3D, self).__init__()
if depth not in self.arch_settings:
raise KeyError('invalid depth {} for resnet'.format(depth))
self.depth = depth
self.pretrained = pretrained
self.num_stages = num_stages
assert num_stages >= 1 and num_stages <= 4
self.spatial_strides = spatial_strides
self.temporal_strides = temporal_strides
self.dilations = dilations
assert len(spatial_strides) == len(temporal_strides) == len(dilations) == num_stages
self.out_indices = out_indices
assert max(out_indices) < num_stages
self.style = style
self.frozen_stages = frozen_stages
self.inflate_freqs = inflate_freq if not isinstance(inflate_freq, int) else (inflate_freq, ) * num_stages
self.bn_eval = bn_eval
self.bn_frozen = bn_frozen
self.partial_bn = partial_bn
self.with_cp = with_cp
self.with_trajectory = with_trajectory
self.trajectory_source_indices = trajectory_source_indices \
if not isinstance(trajectory_source_indices, int) else [trajectory_source_indices, ] * num_stages
self.trajectory_downsample_method = trajectory_downsample_method
self.conv_bias = conv_bias
self.block, stage_blocks = self.arch_settings[depth]
self.stage_blocks = stage_blocks[:num_stages]
for stage in range(num_stages):
self.trajectory_source_indices[stage] = self.trajectory_source_indices[stage] \
if not isinstance(self.trajectory_source_indices[stage], int) else (self.trajectory_source_indices[stage], ) * self.stage_blocks[stage]
self.inplanes = 64
if conv1_kernel_t > 1:
self.conv1 = nn.Conv3d(
3, 64, kernel_size=(1,7,7), stride=(1,2,2), padding=(0,3,3), bias=False)
self.conv1_t = nn.Conv3d(
64, 64, kernel_size=(conv1_kernel_t,1,1), stride=(conv1_stride_t,1,1), padding=((conv1_kernel_t-1)//2,1,1), bias=True)
self.bn1_t = nn.BatchNorm3d(64)
else:
self.conv1 = nn.Conv3d(
3, 64, kernel_size=(1,7,7), stride=(conv1_stride_t,2,2), padding=(0,3,3), bias=False)
self.bn1 = nn.BatchNorm3d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool3d(kernel_size=(pool1_kernel_t,3,3), stride=(pool1_stride_t,2,2), padding=(pool1_kernel_t//2,1,1))
self.use_pool2 = use_pool2
if self.use_pool2:
self.pool2 = nn.MaxPool3d(kernel_size=(3,1,1), stride=(2,1,1), padding=(1,0,0))
self.res_layers = []
for i, num_blocks in enumerate(self.stage_blocks):
traj_src_indices = self.trajectory_source_indices[i] \
if not isinstance(self.trajectory_source_indices[i], int) \
else (self.trajectory_source_indices[i], ) * num_blocks
spatial_stride = spatial_strides[i]
temporal_stride = temporal_strides[i]
dilation = dilations[i]
planes = 64 * 2**i
res_layer = make_res_layer(
self.block,
self.inplanes,
planes,
num_blocks,
spatial_stride=spatial_stride,
temporal_stride=temporal_stride,
dilation=dilation,
style=self.style,
inflate_freq=self.inflate_freqs[i],
with_cp=with_cp,
traj_src_indices=traj_src_indices)
self.inplanes = planes * self.block.expansion
layer_name = 'layer{}'.format(i + 1)
self.add_module(layer_name, res_layer)
self.res_layers.append(layer_name)
self.feat_dim = self.block.expansion * 64 * 2**(
len(self.stage_blocks) - 1)
def __init__(self):
super().__init__()
self.conv0 = nn.Sequential(
nn.Sequential(
nn.Conv3d(1,
32,
kernel_size=(3, 3, 3),
stride=(2, 2, 2),
padding=(1, 1, 1)),
nn.BatchNorm3d(32,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1)),
nn.Sequential(
nn.Conv3d(32,
32,
kernel_size=(3, 3, 3),
stride=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(32,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1)),
nn.Sequential(
nn.Conv3d(32,
32,
kernel_size=(3, 3, 3),
stride=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(32,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1)),
nn.Sequential(
nn.Conv3d(32,
32,
kernel_size=(3, 3, 3),
stride=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(32,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1)))
self.res0 = nn.Conv3d(1,
32,
kernel_size=(3, 3, 3),
stride=(2, 2, 2),
padding=(1, 1, 1))
self.conv1 = nn.Sequential(
nn.Sequential(
nn.Conv3d(32,
64,
kernel_size=(3, 3, 3),
stride=(2, 2, 2),
padding=(1, 1, 1)),
nn.BatchNorm3d(64,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1)),
nn.Sequential(
nn.Conv3d(64,
64,
kernel_size=(3, 3, 3),
stride=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(64,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1)),
nn.Sequential(
nn.Conv3d(64,
64,
kernel_size=(3, 3, 3),
stride=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(64,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1)),
nn.Sequential(
nn.Conv3d(64,
64,
kernel_size=(3, 3, 3),
stride=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(64,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1)))
self.res1 = nn.Conv3d(32,
64,
kernel_size=(3, 3, 3),
stride=(2, 2, 2),
padding=(1, 1, 1))
self.conv2 = nn.Sequential(
nn.Sequential(
nn.Conv3d(64,
128,
kernel_size=(3, 3, 3),
stride=(2, 2, 2),
padding=(1, 1, 1)),
nn.BatchNorm3d(128,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1)),
nn.Sequential(
nn.Conv3d(128,
128,
kernel_size=(3, 3, 3),
stride=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(128,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1)),
nn.Sequential(
nn.Conv3d(128,
128,
kernel_size=(3, 3, 3),
stride=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(128,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1)),
nn.Sequential(
nn.Conv3d(128,
128,
kernel_size=(3, 3, 3),
stride=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(128,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1)))
self.res2 = nn.Conv3d(64,
128,
kernel_size=(3, 3, 3),
stride=(2, 2, 2),
padding=(1, 1, 1))
self.conv3 = nn.Sequential(
nn.Conv3d(128,
256,
kernel_size=(3, 3, 3),
stride=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(256,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1),
nn.Conv3d(256,
256,
kernel_size=(3, 3, 3),
stride=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(256,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1),
nn.Conv3d(256,
256,
kernel_size=(3, 3, 3),
stride=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(256,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1),
nn.Conv3d(256,
256,
kernel_size=(3, 3, 3),
stride=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(256,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1))
self.res3 = nn.Conv3d(128,
256,
kernel_size=(1, 1, 1),
stride=(1, 1, 1))
self.deconv3_0 = nn.Sequential(
nn.ConvTranspose3d(384,
64,
kernel_size=(3, 3, 3),
stride=(2, 2, 2),
padding=(1, 1, 1),
output_padding=(1, 1, 1)),
nn.BatchNorm3d(64,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1))
self.deconv3_1 = nn.Sequential(
nn.Conv3d(64,
64,
kernel_size=(3, 3, 3),
stride=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(64,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1))
self.deconv2_0 = nn.Sequential(
nn.ConvTranspose3d(128,
32,
kernel_size=(3, 3, 3),
stride=(2, 2, 2),
padding=(1, 1, 1),
output_padding=(1, 1, 1)),
nn.BatchNorm3d(32,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1),
)
self.deconv2_1 = nn.Sequential(
nn.Conv3d(32,
32,
kernel_size=(3, 3, 3),
stride=(1, 1, 1),
padding=(1, 1, 1)),
nn.BatchNorm3d(32,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1))
self.deconv1_0 = nn.Sequential(
nn.ConvTranspose3d(64,
6,
kernel_size=(3, 3, 3),
stride=(2, 2, 2),
padding=(1, 1, 1),
output_padding=(1, 1, 1)),
nn.BatchNorm3d(6,
eps=1e-05,
momentum=0.1,
affine=True,
track_running_stats=bn_track),
nn.PReLU(num_parameters=1))
self.deconv1_1 = nn.Conv3d(6,
6,
kernel_size=(3, 3, 3),
stride=(1, 1, 1),
padding=(1, 1, 1))
def __init__(self, args):
super(PSMNet, self).__init__()
self.maxdisp = args.maxdisp
self.planes = args.planes
if args.shuffle:
block3d = ResBlock3DShuffle
block2d = ResBlockShuffle
else:
block3d = ResBlock3D
block2d = ResBlock
self.depth = args.depth
self.sequence = args.sequence
self.flood = args.flood
self.first_conv = nn.Sequential(
nn.Conv2d(3, args.planes*2, kernel_size=3, stride=1, padding=1, dilation=1, bias=False),
nn.BatchNorm2d(args.planes*2),
nn.ReLU(inplace=True),
)
self.first_fuse = nn.Sequential(
nn.Conv2d(args.planes*2, args.planes, kernel_size=1, stride=1, padding=0, dilation=1, bias=False),
nn.BatchNorm2d(args.planes),
nn.ReLU(inplace=True),
)
self.unet_feature = nn.ModuleList()
self.unet_downsample = nn.ModuleList()
self.unet_fuse = nn.ModuleList()
inplanes = args.planes
outplanes = inplanes
for i in range(self.depth):
if i != self.depth - 1:
outplanes = inplanes * 2
self.unet_downsample.append(nn.Sequential(
nn.Conv2d(inplanes, outplanes, kernel_size=3, stride=2, padding=1, dilation=1, bias=False),
nn.BatchNorm2d(outplanes),
nn.ReLU(inplace=True),
))
self.unet_feature.append(nn.Sequential(
block2d(outplanes, kernel_size=3, stride=1, padding=args.dilation, dilation=args.dilation),
block2d(outplanes, kernel_size=3, stride=1, padding=args.dilation, dilation=args.dilation),
))
self.unet_fuse.append(nn.Sequential(
nn.Conv2d(outplanes*2, outplanes, kernel_size=1, stride=1, padding=0, dilation=1, bias=False),
nn.BatchNorm2d(outplanes),
nn.ReLU(inplace=True),
))
inplanes = outplanes
self.cost_merge = nn.Sequential(
nn.Conv3d(outplanes * 2, outplanes, kernel_size=1, stride=1, padding=0, dilation=1, bias=False),
nn.BatchNorm3d(outplanes),
nn.ReLU(inplace=True),
)
self.fusers = nn.ModuleList()
self.classifiers = nn.ModuleList()
self.regressers = nn.ModuleList()
for i in range(self.sequence):
self.fusers.append(nn.Sequential(
nn.Conv3d(outplanes * 2, outplanes, kernel_size=1, stride=1, padding=0, dilation=1, bias=False),
nn.BatchNorm3d(outplanes),
nn.ReLU(inplace=True),
))
self.classifiers.append(nn.Sequential(
block3d(outplanes, kernel_size=3, stride=1, padding=args.dilation, dilation=args.dilation),
block3d(outplanes, kernel_size=3, stride=1, padding=args.dilation, dilation=args.dilation),
))
self.regressers.append(nn.Conv3d(outplanes, 1, kernel_size=1, stride=1, padding=0, dilation=1, bias=False))
self.disparityregression = disparityregression(self.maxdisp//(2**self.depth))
self.refinements = nn.ModuleList()
for i in range(self.depth):
if i != 0:
outplanes = inplanes // 2
self.refinements.append(nn.Sequential(
nn.Conv2d(outplanes*2+1, outplanes, kernel_size=1, stride=1, padding=0, dilation=1, bias=False),
nn.BatchNorm2d(outplanes),
nn.ReLU(inplace=True),
block2d(outplanes, kernel_size=3, stride=1, padding=1, dilation=1),
block2d(outplanes, kernel_size=3, stride=1, padding=1, dilation=1),
nn.Conv2d(outplanes, (2*self.flood + 1)*(2*self.flood + 1),
kernel_size=1, stride=1, padding=0, dilation=1, bias=False),
nn.Sigmoid()
))
inplanes = outplanes
for m in self.modules():
if isinstance(m, nn.Conv2d) or isinstance(m, nn.Conv3d):
torch.nn.init.xavier_normal_(m.weight, gain=1.0)
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.BatchNorm3d):
m.weight.data.fill_(1)
m.bias.data.zero_()
def conv_block_3d(in_dim, out_dim, activation):
return nn.Sequential(
nn.Conv3d(in_dim, out_dim, kernel_size=3, stride=1, padding=1),
nn.BatchNorm3d(out_dim),
activation,
)
