def test_unet2d_encode(num_patches: int,
num_channels: int,
num_output_channels: int,
is_downsampling: bool,
image_shape: TupleInt2) -> None:
"""
Test if the Encode block of a Unet3D correctly works when passing in kernels that only operate in X and Y.
"""
set_random_seed(1234)
layer = UNet3D.UNetEncodeBlock((num_channels, num_output_channels),
kernel_size=(1, 3, 3),
downsampling_stride=(1, 2, 2) if is_downsampling else 1)
input_shape = (num_patches, num_channels) + (1,) + image_shape
input = torch.rand(*input_shape).float()
output = layer(input)
def output_image_size(input_image_size: int) -> int:
# If max pooling is added, it is done with a kernel size of 2, shrinking the image by a factor of 2
image_shrink_factor = 2 if is_downsampling else 1
return input_image_size // image_shrink_factor
# Expected output shape:
# The first dimension (patches) should be retained unchanged.
# We should get as many output channels as requested
# Unet is defined as running over degenerate 3D images with Z=1, this should be preserved.
# The two trailing dimensions are the adjusted image dimensions
expected_output_shape = (num_patches, num_output_channels, 1,
output_image_size(image_shape[0]), output_image_size(image_shape[1]))
assert output.shape == expected_output_shape
def test_unet3d_upsampling() -> None:
assert get_upsampling_kernel_size(2, 3) == (4, 4, 4)
assert get_upsampling_kernel_size(1, 3) == (1, 1, 1)
assert get_upsampling_kernel_size((1, 2, 2), 3) == (1, 4, 4)
assert get_upsampling_kernel_size((1, 2, 3), 3) == (1, 4, 6)
with pytest.raises(ValueError):
get_upsampling_kernel_size((2, 2), 3)
with pytest.raises(ValueError):
get_upsampling_kernel_size(0, 3)
# Test method as integrated into the constructor
assert UNet3D(1, 1, 1, kernel_size=3, downsampling_factor=(1, 2, 2)).upsampling_kernel_size == (1, 4, 4)
def test_unet_summary_generation() -> None:
"""Checks unet summary generation works either in CPU or GPU"""
model = UNet3D(input_image_channels=1,
initial_feature_channels=2,
num_classes=2,
kernel_size=1,
num_downsampling_paths=2)
if machine_has_gpu:
model.cuda()
summary = ModelSummary(model=model).generate_summary(input_sizes=[(1, 4, 4, 4)])
assert summary is not None
def build_net(args: SegmentationModelBase) -> BaseSegmentationModel:
"""
Build network architectures
:param args: Network configuration arguments
"""
full_channels_list = [
args.number_of_image_channels, *args.feature_channels,
args.number_of_classes
]
initial_fcn = [BasicLayer] * 2
residual_blocks = [[BasicLayer, BasicLayer]] * 3
basic_network_definition = initial_fcn + residual_blocks # type: ignore
# no dilation for the initial FCN and then a constant 1 neighbourhood dilation for the rest residual blocks
basic_dilations = [1] * len(initial_fcn) + [2, 2] * len(
basic_network_definition)
# Crop size must be at least 29 because all architectures (apart from UNets) shrink the input image by 28
crop_size_constraints = CropSizeConstraints(
minimum_size=basic_size_shrinkage + 1)
run_weight_initialization = True
network: BaseSegmentationModel
if args.architecture == ModelArchitectureConfig.Basic:
network_definition = basic_network_definition
network = ComplexModel(args, full_channels_list, basic_dilations,
network_definition,
crop_size_constraints) # type: ignore
elif args.architecture == ModelArchitectureConfig.UNet3D:
network = UNet3D(input_image_channels=args.number_of_image_channels,
initial_feature_channels=args.feature_channels[0],
num_classes=args.number_of_classes,
kernel_size=args.kernel_size,
num_downsampling_paths=args.num_downsampling_paths)
run_weight_initialization = False
elif args.architecture == ModelArchitectureConfig.UNet2D:
network = UNet2D(input_image_channels=args.number_of_image_channels,
initial_feature_channels=args.feature_channels[0],
num_classes=args.number_of_classes,
padding_mode=PaddingMode.Edge,
num_downsampling_paths=args.num_downsampling_paths)
run_weight_initialization = False
else:
raise ValueError(f"Unknown model architecture {args.architecture}")
network.validate_crop_size(args.crop_size, "Training crop size")
network.validate_crop_size(args.test_crop_size,
"Test crop size") # type: ignore
# Initialize network weights
if run_weight_initialization:
network.apply(init_weights) # type: ignore
return network
def crop_size_multiple_unet3d(crop_size: TupleInt3,
downsampling_factor: IntOrTuple3,
num_downsampling_paths: int,
expected_crop_size_multiple: TupleInt3) -> None:
model = UNet3D(name="",
input_image_channels=1,
initial_feature_channels=32,
num_classes=2,
kernel_size=3,
downsampling_factor=downsampling_factor,
num_downsampling_paths=num_downsampling_paths)
assert model.crop_size_constraints.multiple_of == expected_crop_size_multiple
model.validate_crop_size(crop_size)
def create_encoder(self, channels: List[int]) -> ModuleList:
"""
Create an image encoder network.
"""
layers = []
for i in range(len(channels) - 1):
layers.append(
UNet3D.UNetEncodeBlock(
channels=(channels[i], channels[i + 1]),
kernel_size=self.kernel_size_per_encoding_block[i],
downsampling_stride=self.stride_size_per_encoding_block[i],
padding_mode=self.padding_mode,
use_residual=False,
depth=i,
))
return ModuleList(layers)
def test_unet_model_parallel() -> None:
"""Checks model parallel utilises all the available GPU devices for forward pass"""
if no_gpu_available:
logging.warning("CUDA capable GPU is not available - UNet Model Parallel cannot be tested")
return
model = UNet3D(input_image_channels=1,
initial_feature_channels=2,
num_classes=2,
kernel_size=1,
num_downsampling_paths=2).cuda()
# Partition the network across all available gpu
available_devices = [torch.device('cuda:{}'.format(ii)) for ii in range(torch.cuda.device_count())]
model.generate_model_summary()
model.partition_model(devices=available_devices)
# Verify that all the devices are utilised by the layers of the model
summary = ModelSummary(model=model).generate_summary(input_sizes=[(1, 4, 4, 4)])
layer_devices = set()
for layer_summary in summary.values():
if layer_summary.device:
layer_devices.add(layer_summary.device)
assert layer_devices == set(available_devices)