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