def plan_pool_and_conv_pool_late(
self,
patch_size: Sequence[int],
spacing: Sequence[float],
) -> Tuple[List[int], List[Tuple[int]], List[Tuple[int]], Sequence[int],
Sequence[int]]:
"""
Plan pooling and convolutions of encoder network
Axis which do not need pooling in every block are pooled as late as possible
Uses kernel size 1 for anisotropic axis which are not reached by the fov yet
Args:
patch_size: target path size
spacing: target spacing transposed
Returns:
List[int]: max number of pooling operations per axis
List[Tuple[int]]: kernel sizes of pooling operations
List[Tuple[int]]: kernel sizes of convolution layers
Sequence[int]: patch size
Sequence[int]: coefficient each axes needs to be divisable by
"""
num_pool_per_axis, pool_op_kernel_sizes, conv_kernel_sizes, \
patch_size, must_be_divisible_by = get_pool_and_conv_props(
spacing=spacing, patch_size=patch_size,
min_feature_map_size=self.min_feature_map_size,
max_numpool=self.max_num_pool)
return num_pool_per_axis, pool_op_kernel_sizes, conv_kernel_sizes, patch_size, must_be_divisible_by
def get_properties_for_stage(self, current_spacing, original_spacing,
original_shape, num_cases, num_modalities,
num_classes):
new_median_shape = np.round(original_spacing / current_spacing *
original_shape).astype(int)
dataset_num_voxels = np.prod(new_median_shape,
dtype=np.int64) * num_cases
input_patch_size = new_median_shape[1:]
network_numpool, net_pool_kernel_sizes, net_conv_kernel_sizes, input_patch_size, \
shape_must_be_divisible_by = get_pool_and_conv_props(current_spacing[1:], input_patch_size,
self.unet_featuremap_min_edge_length,
self.unet_max_numpool)
# we pretend to use 30 feature maps. This will yield the same configuration as in V1. The larger memory
# footpring of 32 vs 30 is mor ethan offset by the fp16 training. We make fp16 training default
# Reason for 32 vs 30 feature maps is that 32 is faster in fp16 training (because multiple of 8)
estimated_gpu_ram_consumption = Generic_UNet.compute_approx_vram_consumption(
input_patch_size,
network_numpool,
30,
self.unet_max_num_filters,
num_modalities,
num_classes,
net_pool_kernel_sizes,
conv_per_stage=self.conv_per_stage)
batch_size = int(
np.floor(Generic_UNet.use_this_for_batch_size_computation_2D /
estimated_gpu_ram_consumption *
Generic_UNet.DEFAULT_BATCH_SIZE_2D))
if batch_size < self.unet_min_batch_size:
raise RuntimeError(
"This framework is not made to process patches this large. We will add patch-based "
"2D networks later. Sorry for the inconvenience")
# check if batch size is too large (more than 5 % of dataset)
max_batch_size = np.round(
self.batch_size_covers_max_percent_of_dataset *
dataset_num_voxels /
np.prod(input_patch_size, dtype=np.int64)).astype(int)
batch_size = min(batch_size, max_batch_size)
plan = {
'batch_size': batch_size,
'num_pool_per_axis': network_numpool,
'patch_size': input_patch_size,
'median_patient_size_in_voxels': new_median_shape,
'current_spacing': current_spacing,
'original_spacing': original_spacing,
'pool_op_kernel_sizes': net_pool_kernel_sizes,
'conv_kernel_sizes': net_conv_kernel_sizes,
'do_dummy_2D_data_aug': False
}
return plan
def get_properties_for_stage(current_spacing, original_spacing,
original_shape, num_cases, num_modalities,
num_classes, transpose_forward):
spacing_transposed = current_spacing[transpose_forward]
new_median_shape = np.round(original_spacing / current_spacing *
original_shape).astype(int)
new_median_shape_transposed = new_median_shape[transpose_forward]
dataset_num_voxels = np.prod(new_median_shape) * num_cases
input_patch_size = new_median_shape_transposed[1:]
network_numpool, net_pool_kernel_sizes, net_conv_kernel_sizes, input_patch_size, \
shape_must_be_divisible_by = get_pool_and_conv_props(spacing_transposed[1:], input_patch_size,
FEATUREMAP_MIN_EDGE_LENGTH_BOTTLENECK,
Generic_UNet.MAX_NUMPOOL_2D)
estimated_gpu_ram_consumption = Generic_UNet.compute_approx_vram_consumption(
input_patch_size, network_numpool,
Generic_UNet.BASE_NUM_FEATURES_2D, Generic_UNet.MAX_FILTERS_2D,
num_modalities, num_classes, net_pool_kernel_sizes)
batch_size = int(
np.floor(Generic_UNet.use_this_for_batch_size_computation_2D /
estimated_gpu_ram_consumption *
Generic_UNet.DEFAULT_BATCH_SIZE_2D))
if batch_size < dataset_min_batch_size_cap:
raise RuntimeError(
"This framework is not made to process patches this large. We will add patch-based "
"2D networks later. Sorry for the inconvenience")
# check if batch size is too large (more than 5 % of dataset)
max_batch_size = np.round(
batch_size_covers_max_percent_of_dataset * dataset_num_voxels /
np.prod(input_patch_size)).astype(int)
batch_size = min(batch_size, max_batch_size)
plan = {
'batch_size': batch_size,
'num_pool_per_axis': network_numpool,
'patch_size': input_patch_size,
'median_patient_size_in_voxels': new_median_shape,
'current_spacing': current_spacing,
'original_spacing': original_spacing,
'pool_op_kernel_sizes': net_pool_kernel_sizes,
'conv_kernel_sizes': net_conv_kernel_sizes,
'do_dummy_2D_data_aug': False
}
return plan
def get_properties_for_stage(self, current_spacing, original_spacing,
original_shape, num_cases, num_modalities,
num_classes):
"""
We need to adapt ref
"""
new_median_shape = np.round(original_spacing / current_spacing *
original_shape).astype(int)
dataset_num_voxels = np.prod(new_median_shape) * num_cases
# the next line is what we had before as a default. The patch size had the same aspect ratio as the median shape of a patient. We swapped t
# input_patch_size = new_median_shape
# compute how many voxels are one mm
input_patch_size = 1 / np.array(current_spacing)
# normalize voxels per mm
input_patch_size /= input_patch_size.mean()
# create an isotropic patch of size 512x512x512mm
input_patch_size *= 1 / min(
input_patch_size) * 512 # to get a starting value
input_patch_size = np.round(input_patch_size).astype(int)
# clip it to the median shape of the dataset because patches larger then that make not much sense
input_patch_size = [
min(i, j) for i, j in zip(input_patch_size, new_median_shape)
]
network_num_pool_per_axis, pool_op_kernel_sizes, conv_kernel_sizes, new_shp, \
shape_must_be_divisible_by = get_pool_and_conv_props(current_spacing, input_patch_size,
self.unet_featuremap_min_edge_length,
self.unet_max_numpool)
# use_this_for_batch_size_computation_3D = 520000000 # 505789440
# typical ExperimentPlanner3D_v21 configurations use 7.5GB, but on a V100 we have 32. Allow for more space
# to be used
ref = Generic_UNet.use_this_for_batch_size_computation_3D * 32 / 8
here = Generic_UNet.compute_approx_vram_consumption(
new_shp,
network_num_pool_per_axis,
self.unet_base_num_features,
self.unet_max_num_filters,
num_modalities,
num_classes,
pool_op_kernel_sizes,
conv_per_stage=self.conv_per_stage)
while here > ref:
axis_to_be_reduced = np.argsort(new_shp / new_median_shape)[-1]
tmp = deepcopy(new_shp)
tmp[axis_to_be_reduced] -= shape_must_be_divisible_by[
axis_to_be_reduced]
_, _, _, _, shape_must_be_divisible_by_new = \
get_pool_and_conv_props(current_spacing, tmp,
self.unet_featuremap_min_edge_length,
self.unet_max_numpool,
)
new_shp[axis_to_be_reduced] -= shape_must_be_divisible_by_new[
axis_to_be_reduced]
# we have to recompute numpool now:
network_num_pool_per_axis, pool_op_kernel_sizes, conv_kernel_sizes, new_shp, \
shape_must_be_divisible_by = get_pool_and_conv_props(current_spacing, new_shp,
self.unet_featuremap_min_edge_length,
self.unet_max_numpool,
)
here = Generic_UNet.compute_approx_vram_consumption(
new_shp,
network_num_pool_per_axis,
self.unet_base_num_features,
self.unet_max_num_filters,
num_modalities,
num_classes,
pool_op_kernel_sizes,
conv_per_stage=self.conv_per_stage)
# print(new_shp)
input_patch_size = new_shp
batch_size = Generic_UNet.DEFAULT_BATCH_SIZE_3D # This is what wirks with 128**3
batch_size = int(np.floor(max(ref / here, 1) * batch_size))
# check if batch size is too large
max_batch_size = np.round(
self.batch_size_covers_max_percent_of_dataset *
dataset_num_voxels /
np.prod(input_patch_size, dtype=np.int64)).astype(int)
max_batch_size = max(max_batch_size, self.unet_min_batch_size)
batch_size = max(1, min(batch_size, max_batch_size))
do_dummy_2D_data_aug = (max(input_patch_size) / input_patch_size[0]
) > self.anisotropy_threshold
plan = {
'batch_size': batch_size,
'num_pool_per_axis': network_num_pool_per_axis,
'patch_size': input_patch_size,
'median_patient_size_in_voxels': new_median_shape,
'current_spacing': current_spacing,
'original_spacing': original_spacing,
'do_dummy_2D_data_aug': do_dummy_2D_data_aug,
'pool_op_kernel_sizes': pool_op_kernel_sizes,
'conv_kernel_sizes': conv_kernel_sizes,
}
return plan
def get_properties_for_stage(self, current_spacing, original_spacing,
original_shape, num_cases, num_modalities,
num_classes):
"""
ExperimentPlanner configures pooling so that we pool late. Meaning that if the number of pooling per axis is
(2, 3, 3), then the first pooling operation will always pool axes 1 and 2 and not 0, irrespective of spacing.
This can cause a larger memory footprint, so it can be beneficial to revise this.
Here we are pooling based on the spacing of the data.
"""
new_median_shape = np.round(original_spacing / current_spacing *
original_shape).astype(int)
dataset_num_voxels = np.prod(new_median_shape) * num_cases
# the next line is what we had before as a default. The patch size had the same aspect ratio as the median shape of a patient. We swapped t
# input_patch_size = new_median_shape
# compute how many voxels are one mm
input_patch_size = 1 / np.array(current_spacing)
# normalize voxels per mm
input_patch_size /= input_patch_size.mean()
# create an isotropic patch of size 512x512x512mm
input_patch_size *= 1 / min(
input_patch_size) * 512 # to get a starting value
input_patch_size = np.round(input_patch_size).astype(int)
# clip it to the median shape of the dataset because patches larger then that make not much sense
input_patch_size = [
min(i, j) for i, j in zip(input_patch_size, new_median_shape)
]
network_num_pool_per_axis, pool_op_kernel_sizes, conv_kernel_sizes, new_shp, \
shape_must_be_divisible_by = get_pool_and_conv_props(current_spacing, input_patch_size,
self.unet_featuremap_min_edge_length,
self.unet_max_numpool)
# we compute as if we were using only 30 feature maps. We can do that because fp16 training is the standard
# now. That frees up some space. The decision to go with 32 is solely due to the speedup we get (non-multiples
# of 8 are not supported in nvidia amp)
ref = ResAxialAttentionUNet.use_this_for_batch_size_computation_3D * self.unet_base_num_features / \
ResAxialAttentionUNet.BASE_NUM_FEATURES_3D
here = ResAxialAttentionUNet.compute_approx_vram_consumption(
new_shp,
network_num_pool_per_axis,
self.unet_base_num_features,
self.unet_max_num_filters,
num_modalities,
num_classes,
pool_op_kernel_sizes,
conv_per_stage=self.conv_per_stage)
while here > ref:
axis_to_be_reduced = np.argsort(new_shp / new_median_shape)[-1]
tmp = deepcopy(new_shp)
tmp[axis_to_be_reduced] -= shape_must_be_divisible_by[
axis_to_be_reduced]
_, _, _, _, shape_must_be_divisible_by_new = \
get_pool_and_conv_props(current_spacing, tmp,
self.unet_featuremap_min_edge_length,
self.unet_max_numpool,
)
new_shp[axis_to_be_reduced] -= shape_must_be_divisible_by_new[
axis_to_be_reduced]
# we have to recompute numpool now:
network_num_pool_per_axis, pool_op_kernel_sizes, conv_kernel_sizes, new_shp, \
shape_must_be_divisible_by = get_pool_and_conv_props(current_spacing, new_shp,
self.unet_featuremap_min_edge_length,
self.unet_max_numpool,
)
here = Generic_UNet.compute_approx_vram_consumption(
new_shp,
network_num_pool_per_axis,
self.unet_base_num_features,
self.unet_max_num_filters,
num_modalities,
num_classes,
pool_op_kernel_sizes,
conv_per_stage=self.conv_per_stage)
#print(new_shp)
#print(here, ref)
input_patch_size = new_shp
batch_size = Generic_UNet.DEFAULT_BATCH_SIZE_3D # This is what wirks with 128**3
batch_size = int(np.floor(max(ref / here, 1) * batch_size))
# check if batch size is too large
max_batch_size = np.round(
self.batch_size_covers_max_percent_of_dataset *
dataset_num_voxels /
np.prod(input_patch_size, dtype=np.int64)).astype(int)
max_batch_size = max(max_batch_size, self.unet_min_batch_size)
batch_size = min(batch_size, max_batch_size)
do_dummy_2D_data_aug = (max(input_patch_size) / input_patch_size[0]
) > self.anisotropy_threshold
plan = {
#liuyiyao
#'batch_size': batch_size,
'batch_size': 4,
'num_pool_per_axis': network_num_pool_per_axis,
#'patch_size': input_patch_size,
#liuyiyao
'patch_size': (8, 128, 128),
'median_patient_size_in_voxels': new_median_shape,
'current_spacing': current_spacing,
'original_spacing': original_spacing,
'do_dummy_2D_data_aug': do_dummy_2D_data_aug,
'pool_op_kernel_sizes': pool_op_kernel_sizes,
'conv_kernel_sizes': conv_kernel_sizes,
}
return plan
def get_properties_for_stage(self, current_spacing, original_spacing,
original_shape, num_cases, num_modalities,
num_classes):
"""
We use FabiansUNet instead of Generic_UNet
"""
new_median_shape = np.round(original_spacing / current_spacing *
original_shape).astype(int)
dataset_num_voxels = np.prod(new_median_shape) * num_cases
# the next line is what we had before as a default. The patch size had the same aspect ratio as the median shape of a patient. We swapped t
# input_patch_size = new_median_shape
# compute how many voxels are one mm
input_patch_size = 1 / np.array(current_spacing)
# normalize voxels per mm
input_patch_size /= input_patch_size.mean()
# create an isotropic patch of size 512x512x512mm
input_patch_size *= 1 / min(
input_patch_size) * 512 # to get a starting value
input_patch_size = np.round(input_patch_size).astype(int)
# clip it to the median shape of the dataset because patches larger then that make not much sense
input_patch_size = [
min(i, j) for i, j in zip(input_patch_size, new_median_shape)
]
network_num_pool_per_axis, pool_op_kernel_sizes, conv_kernel_sizes, new_shp, \
shape_must_be_divisible_by = get_pool_and_conv_props(current_spacing, input_patch_size,
self.unet_featuremap_min_edge_length,
self.unet_max_numpool)
pool_op_kernel_sizes = [[1, 1, 1]] + pool_op_kernel_sizes
blocks_per_stage_encoder = FabiansUNet.default_blocks_per_stage_encoder[:len(
pool_op_kernel_sizes)]
blocks_per_stage_decoder = FabiansUNet.default_blocks_per_stage_decoder[:len(
pool_op_kernel_sizes) - 1]
ref = FabiansUNet.use_this_for_3D_configuration
here = FabiansUNet.compute_approx_vram_consumption(
input_patch_size,
self.unet_base_num_features,
self.unet_max_num_filters,
num_modalities,
num_classes,
pool_op_kernel_sizes,
blocks_per_stage_encoder,
blocks_per_stage_decoder,
2,
self.unet_min_batch_size,
)
while here > ref:
axis_to_be_reduced = np.argsort(new_shp / new_median_shape)[-1]
tmp = deepcopy(new_shp)
tmp[axis_to_be_reduced] -= shape_must_be_divisible_by[
axis_to_be_reduced]
_, _, _, _, shape_must_be_divisible_by_new = \
get_pool_and_conv_props(current_spacing, tmp,
self.unet_featuremap_min_edge_length,
self.unet_max_numpool,
)
new_shp[axis_to_be_reduced] -= shape_must_be_divisible_by_new[
axis_to_be_reduced]
# we have to recompute numpool now:
network_num_pool_per_axis, pool_op_kernel_sizes, conv_kernel_sizes, new_shp, \
shape_must_be_divisible_by = get_pool_and_conv_props(current_spacing, new_shp,
self.unet_featuremap_min_edge_length,
self.unet_max_numpool,
)
pool_op_kernel_sizes = [[1, 1, 1]] + pool_op_kernel_sizes
blocks_per_stage_encoder = FabiansUNet.default_blocks_per_stage_encoder[:len(
pool_op_kernel_sizes)]
blocks_per_stage_decoder = FabiansUNet.default_blocks_per_stage_decoder[:len(
pool_op_kernel_sizes) - 1]
here = FabiansUNet.compute_approx_vram_consumption(
new_shp, self.unet_base_num_features,
self.unet_max_num_filters, num_modalities, num_classes,
pool_op_kernel_sizes, blocks_per_stage_encoder,
blocks_per_stage_decoder, 2, self.unet_min_batch_size)
input_patch_size = new_shp
batch_size = FabiansUNet.default_min_batch_size
batch_size = int(np.floor(max(ref / here, 1) * batch_size))
# check if batch size is too large
max_batch_size = np.round(
self.batch_size_covers_max_percent_of_dataset *
dataset_num_voxels /
np.prod(input_patch_size, dtype=np.int64)).astype(int)
max_batch_size = max(max_batch_size, self.unet_min_batch_size)
batch_size = min(batch_size, max_batch_size)
do_dummy_2D_data_aug = (max(input_patch_size) / input_patch_size[0]
) > self.anisotropy_threshold
plan = {
'batch_size': batch_size,
'num_pool_per_axis': network_num_pool_per_axis,
'patch_size': input_patch_size,
'median_patient_size_in_voxels': new_median_shape,
'current_spacing': current_spacing,
'original_spacing': original_spacing,
'do_dummy_2D_data_aug': do_dummy_2D_data_aug,
'pool_op_kernel_sizes': pool_op_kernel_sizes,
'conv_kernel_sizes': conv_kernel_sizes,
'num_blocks_encoder': blocks_per_stage_encoder,
'num_blocks_decoder': blocks_per_stage_decoder
}
return plan
def get_properties_for_stage(self, current_spacing, original_spacing,
original_shape, num_cases, num_modalities,
num_classes):
new_median_shape = np.round(original_spacing / current_spacing *
original_shape).astype(int)
dataset_num_voxels = np.prod(new_median_shape,
dtype=np.int64) * num_cases
input_patch_size = new_median_shape[1:]
network_num_pool_per_axis, pool_op_kernel_sizes, conv_kernel_sizes, new_shp, \
shape_must_be_divisible_by = get_pool_and_conv_props(current_spacing[1:], input_patch_size,
self.unet_featuremap_min_edge_length,
self.unet_max_numpool)
# we pretend to use 30 feature maps. This will yield the same configuration as in V1. The larger memory
# footpring of 32 vs 30 is mor ethan offset by the fp16 training. We make fp16 training default
# Reason for 32 vs 30 feature maps is that 32 is faster in fp16 training (because multiple of 8)
ref = Generic_UNet.use_this_for_batch_size_computation_2D * Generic_UNet.DEFAULT_BATCH_SIZE_2D / 2 # for batch size 2
here = Generic_UNet.compute_approx_vram_consumption(
new_shp,
network_num_pool_per_axis,
30,
self.unet_max_num_filters,
num_modalities,
num_classes,
pool_op_kernel_sizes,
conv_per_stage=self.conv_per_stage)
while here > ref:
axis_to_be_reduced = np.argsort(new_shp / new_median_shape[1:])[-1]
tmp = deepcopy(new_shp)
tmp[axis_to_be_reduced] -= shape_must_be_divisible_by[
axis_to_be_reduced]
_, _, _, _, shape_must_be_divisible_by_new = \
get_pool_and_conv_props(current_spacing[1:], tmp, self.unet_featuremap_min_edge_length,
self.unet_max_numpool)
new_shp[axis_to_be_reduced] -= shape_must_be_divisible_by_new[
axis_to_be_reduced]
# we have to recompute numpool now:
network_num_pool_per_axis, pool_op_kernel_sizes, conv_kernel_sizes, new_shp, \
shape_must_be_divisible_by = get_pool_and_conv_props(current_spacing[1:], new_shp,
self.unet_featuremap_min_edge_length,
self.unet_max_numpool)
here = Generic_UNet.compute_approx_vram_consumption(
new_shp,
network_num_pool_per_axis,
self.unet_base_num_features,
self.unet_max_num_filters,
num_modalities,
num_classes,
pool_op_kernel_sizes,
conv_per_stage=self.conv_per_stage)
# print(new_shp)
batch_size = int(np.floor(ref / here) * 2)
input_patch_size = new_shp
if batch_size < self.unet_min_batch_size:
raise RuntimeError("This should not happen")
# check if batch size is too large (more than 5 % of dataset)
max_batch_size = np.round(
self.batch_size_covers_max_percent_of_dataset *
dataset_num_voxels /
np.prod(input_patch_size, dtype=np.int64)).astype(int)
batch_size = max(1, min(batch_size, max_batch_size))
plan = {
'batch_size': batch_size,
'num_pool_per_axis': network_num_pool_per_axis,
'patch_size': input_patch_size,
'median_patient_size_in_voxels': new_median_shape,
'current_spacing': current_spacing,
'original_spacing': original_spacing,
'pool_op_kernel_sizes': pool_op_kernel_sizes,
'conv_kernel_sizes': conv_kernel_sizes,
'do_dummy_2D_data_aug': False
}
return plan
