def test_ResNet(MCdataset):
datagen = DataGen(MCdataset, LabelType.Classification)
data = [datagen[i] for i in range(3)]
coord, ener, lab, ev = collatefn(data)
spatial_size = (51, 51, 51)
init_conv_nplanes = 4
init_conv_kernel = 3
kernel_sizes = [7, 5, 3]
stride_sizes = [2, 2]
basic_num = 3
nclasses = 2
nlinear = 8
net = ResNet(spatial_size,
init_conv_nplanes,
init_conv_kernel,
kernel_sizes,
stride_sizes,
basic_num,
nlinear=nlinear,
nclasses=nclasses)
out = net((coord, ener))
assert out.shape[0] == len(lab)
assert out.shape[1] == nclasses
def test_ResidualBlock_basic(MCdataset):
datagen = DataGen(MCdataset, LabelType.Classification)
data = [datagen[i] for i in range(3)]
coord, ener, lab, ev = collatefn(data)
spatial_size = (50, 50, 50)
dim = 3
inplanes = 1
kernel = 2
x = scn.InputLayer(dim, spatial_size)((coord, ener))
out = ResidualBlock_basic(inplanes, kernel)(x)
assert out.features.shape[1] == inplanes
for i, size in enumerate(spatial_size):
assert out.spatial_size[i] == size
def test_ConvBNBlock(MCdataset):
datagen = DataGen(MCdataset, LabelType.Classification)
data = [datagen[i] for i in range(3)]
coord, ener, lab, ev = collatefn(data)
spatial_size = (50, 50, 50)
dim = 3
inplanes = 1
outplanes = 3
kernel = 2
stride = 2
x = scn.InputLayer(dim, spatial_size)((coord, ener))
out = ConvBNBlock(inplanes, outplanes, kernel)(x)
out_with_stride = ConvBNBlock(inplanes, outplanes, kernel,
stride=stride)(x)
for i, size in enumerate(spatial_size):
assert out.spatial_size[i] == size
assert out_with_stride.spatial_size[i] == (size - kernel) / stride + 1
def test_ResidualBlock_upsample(MCdataset):
datagen = DataGen(MCdataset, LabelType.Classification)
data = [datagen[i] for i in range(3)]
coord, ener, lab, ev = collatefn(data)
spatial_size = (50, 50, 50)
dim = 3
inplanes = 1
outplanes = 6
kernel = 2
stride = 2
x = scn.InputLayer(dim, spatial_size)((coord, ener))
x = scn.SubmanifoldConvolution(dim, inplanes, outplanes, kernel, False)(x)
inplanes = x.features.shape[1]
out = ResidualBlock_upsample(inplanes, kernel, stride)(x)
assert out.features.shape[1] == inplanes / 2
for i, size in enumerate(spatial_size):
assert out.spatial_size[i] == kernel + stride * (size - 1)
def test_UNet(MCdataset):
datagen = DataGen(MCdataset, LabelType.Classification)
data = [datagen[i] for i in range(3)]
coord, ener, lab, ev = collatefn(data)
spatial_size = (51, 51, 51)
init_conv_nplanes = 4
init_conv_kernel = 3
kernel_sizes = [7, 5, 3]
stride_sizes = [2, 2]
basic_num = 3
nclasses = 3
net = UNet(spatial_size,
init_conv_nplanes,
init_conv_kernel,
kernel_sizes,
stride_sizes,
basic_num,
nclasses=nclasses)
last_basic = []
net.basic_up[0][2].add.register_forward_hook(
lambda model, input, output: last_basic.append(
[output.spatial_size, output.features.shape]))
assert len(net.downsample) == len(kernel_sizes) - 1
assert len(net.upsample) == len(kernel_sizes) - 1
assert len(net.basic_down) == len(kernel_sizes) - 1
assert len(net.basic_up) == len(kernel_sizes) - 1
assert len(net.basic_down[0]) == basic_num
out = net.forward((coord, ener))
for i, size in enumerate(last_basic[0][0]):
assert size == spatial_size[i]
assert last_basic[0][1][1] == init_conv_nplanes
assert out.size()[0] == coord.size()[0]
assert out.size()[1] == nclasses