def conv_tr(in_channels,
out_channels,
kernel_size,
stride=1,
dilation=1,
bias=False,
region_type=ME.RegionType.HYPER_CUBE,
dimension=-1):
assert dimension > 0, 'Dimension must be a positive integer'
kernel_generator = ME.KernelGenerator(kernel_size,
stride,
dilation,
is_transpose=True,
region_type=region_type,
dimension=dimension)
kernel_generator = ME.KernelGenerator(kernel_size,
stride,
dilation,
is_transpose=True,
region_type=region_type,
dimension=dimension)
return ME.MinkowskiConvolutionTranspose(in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
dilation=dilation,
bias=bias,
kernel_generator=kernel_generator,
dimension=dimension)
def __init__(self, inc, outc, ks=3, stride=1, dilation=1, D=4):
super(ResidualBlock4d, self).__init__()
self.net = nn.Sequential(
ME.MinkowskiConvolution(
inc, outc, dimension=D,
kernel_generator=ME.KernelGenerator(
kernel_size=ks, dimension=D,
region_type=ME.RegionType.HYBRID,
axis_types=(ME.RegionType.HYPERCUBE, ME.RegionType.HYPERCUBE,
ME.RegionType.HYPERCUBE, ME.RegionType.HYPERCROSS))),
ME.MinkowskiBatchNorm(outc),
ME.MinkowskiReLU(True),
ME.MinkowskiConvolution(
outc, outc, dimension=D,
kernel_generator=ME.KernelGenerator(
kernel_size=ks, dimension=D,
region_type=ME.RegionType.HYBRID,
axis_types=(ME.RegionType.HYPERCUBE, ME.RegionType.HYPERCUBE,
ME.RegionType.HYPERCUBE, ME.RegionType.HYPERCROSS))),
ME.MinkowskiBatchNorm(outc))
nn.init.constant_(self.net[1].bn.weight, 1.0)
nn.init.constant_(self.net[1].bn.bias, 0.0)
nn.init.constant_(self.net[4].bn.weight, 1.0)
nn.init.constant_(self.net[4].bn.bias, 0.0)
self.downsample = nn.Sequential() if (inc == outc and stride == 1) else nn.Sequential(
ME.MinkowskiConvolution(
inc, outc, kernel_size=1, dilation=1, stride=stride, dimension=D),
ME.MinkowskiBatchNorm(outc))
if len(self.downsample) > 0:
nn.init.constant_(self.downsample[1].bn.weight, 1.0)
nn.init.constant_(self.downsample[1].bn.bias, 0.0)
self.relu = ME.MinkowskiReLU(True)
def conv_tr(in_planes,
out_planes,
kernel_size,
upsample_stride=1,
dilation=1,
bias=False,
conv_type=ConvType.HYPERCUBE,
D=-1):
assert D > 0, 'Dimension must be a positive integer'
region_type, axis_types, kernel_size = convert_conv_type(conv_type, kernel_size, D)
kernel_generator = ME.KernelGenerator(
kernel_size,
upsample_stride,
dilation,
region_type=region_type,
axis_types=axis_types,
dimension=D)
return ME.MinkowskiConvolutionTranspose(
in_channels=in_planes,
out_channels=out_planes,
kernel_size=kernel_size,
stride=upsample_stride,
dilation=dilation,
has_bias=bias,
kernel_generator=kernel_generator,
dimension=D)
def conv(in_channels,
out_channels,
kernel_size=3,
stride=1,
dilation=1,
bias=False,
region_type=0,
dimension=3):
if not isinstance(region_type, ME.RegionType):
if region_type == 0:
region_type = ME.RegionType.HYPER_CUBE
elif region_type == 1:
region_type = ME.RegionType.HYPER_CROSS
else:
raise ValueError('Unsupported region type')
kernel_generator = ME.KernelGenerator(kernel_size=kernel_size,
stride=stride,
dilation=dilation,
region_type=region_type,
dimension=dimension)
return ME.MinkowskiConvolution(in_channels,
out_channels,
kernel_size=kernel_size,
stride=stride,
kernel_generator=kernel_generator,
dimension=dimension)
def sum_pool(kernel_size, stride=1, dilation=1, conv_type=ConvType.HYPERCUBE, D=-1):
assert D > 0, 'Dimension must be a positive integer'
region_type, axis_types, kernel_size = convert_conv_type(conv_type, kernel_size, D)
kernel_generator = ME.KernelGenerator(
kernel_size, stride, dilation, region_type=region_type, axis_types=axis_types, dimension=D)
return ME.MinkowskiSumPooling(
kernel_size=kernel_size,
stride=stride,
dilation=dilation,
kernel_generator=kernel_generator,
dimension=D)
def __init__(self, inc, outc, ks=3, stride=1, dilation=1, D=3):
super(BasicConvolutionBlock4d, self).__init__()
self.net = nn.Sequential(
ME.MinkowskiConvolution(
inc, outc, dilation=dilation, stride=stride, dimension=D,
kernel_generator=ME.KernelGenerator(
kernel_size=(ks, ks, ks, 1), dimension=D,
region_type=ME.RegionType.HYPERCUBE)),
ME.MinkowskiBatchNorm(outc),
ME.MinkowskiReLU(True))
nn.init.constant_(self.net[1].bn.weight, 1.0)
nn.init.constant_(self.net[1].bn.bias, 0.0)
"dilation": 2,
# "region_type":ME.RegionType.HYPER_CROSS,
"region_type": ME.RegionType.CUSTOM,
"region_offsets": ro0,
"dimension": 3,
}
print(ro0)
# import ipdb; ipdb.set_trace()
kg0 = ME.KernelGenerator(
# **kgargs
kernel_size=2,
stride=1,
dilation=2,
region_type=ME.RegionType.CUSTOM,
region_offsets=ro0,
dimension=3
)
conv = ME.MinkowskiChannelwiseConvolution(
in_channels=1,
kernel_size=2,
stride=1,
dilation=1,
bias=False,
kernel_generator=kg0,
dimension=3
)
def __init__(self, cfg):
super(MinkUNet4d, self).__init__()
self.stem = nn.Sequential(
ME.MinkowskiConvolution(
3, 32, dimension=4,
kernel_generator=ME.KernelGenerator(
kernel_size=(5, 5, 5, 1), dimension=4,
region_type=ME.RegionType.HYPERCUBE)),
ME.MinkowskiBatchNorm(32),
ME.MinkowskiReLU(True))
nn.init.constant_(self.stem[1].bn.weight, 1.0)
nn.init.constant_(self.stem[1].bn.bias, 0.0)
self.stage1 = nn.Sequential(
BasicConvolutionBlock4d(32, 32, ks=2, stride=(2, 2, 2, 1), D=4),
ResidualBlock4d(32, 32, ks=3),
ResidualBlock4d(32, 32, ks=3),
)
self.stage2 = nn.Sequential(
BasicConvolutionBlock4d(32, 32, ks=2, stride=(2, 2, 2, 1), D=4),
ResidualBlock4d(32, 64, ks=3),
ResidualBlock4d(64, 64, ks=3)
)
self.stage3 = nn.Sequential(
BasicConvolutionBlock4d(64, 64, ks=2, stride=(2, 2, 2, 1), D=4),
ResidualBlock4d(64, 128, ks=3),
ResidualBlock4d(128, 128, ks=3),
)
self.stage4 = nn.Sequential(
BasicConvolutionBlock4d(128, 128, ks=2, stride=(2, 2, 2, 1), D=4),
ResidualBlock4d(128, 256, ks=3),
ResidualBlock4d(256, 256, ks=3),
)
self.up1 = nn.ModuleList([
BasicDeconvolutionBlock4d(
256, 256, ks=2, stride=(2, 2, 2, 1), D=4),
nn.Sequential(
ResidualBlock4d(384, 256, ks=3),
ResidualBlock4d(256, 256, ks=3),
)
])
self.up2 = nn.ModuleList([
BasicDeconvolutionBlock4d(
256, 128, ks=2, stride=(2, 2, 2, 1), D=4),
nn.Sequential(
ResidualBlock4d(192, 128, ks=3),
ResidualBlock4d(128, 128, ks=3),
)
])
self.up3 = nn.ModuleList([
BasicDeconvolutionBlock4d(128, 96, ks=2, stride=(2, 2, 2, 1), D=4),
nn.Sequential(
ResidualBlock4d(128, 96, ks=3),
ResidualBlock4d(96, 96, ks=3),
)
])
self.up4 = nn.ModuleList([
BasicDeconvolutionBlock4d(96, 96, ks=2, stride=(2, 2, 2, 1), D=4),
nn.Sequential(
ResidualBlock4d(128, 96, ks=3),
ResidualBlock4d(96, 96, ks=3),
)
])
self.classifier = nn.Sequential(
ME.MinkowskiConvolution(
96, cfg.MODEL.NUM_CLASSES,
kernel_size=1, dimension=4))
