| import logging |
| import os |
|
|
| import numpy as np |
| import torch |
| import torch.distributed as dist |
|
|
| from monai.apps.auto3dseg.auto_runner import logger |
|
|
| print = logger.debug |
| roi_size_default = [224, 224, 144] |
|
|
|
|
| def logger_configure(log_output_file: str = None, debug=False, global_rank=0) -> None: |
| log_config = { |
| "version": 1, |
| "disable_existing_loggers": False, |
| "formatters": {"monai_default": {"format": "%(message)s"}}, |
| "loggers": { |
| "monai.apps.auto3dseg.auto_runner": {"handlers": ["console", "file"], "level": "DEBUG", "propagate": False} |
| }, |
| |
| "handlers": { |
| "console": { |
| "class": "logging.StreamHandler", |
| "level": "INFO", |
| "formatter": "monai_default", |
| |
| }, |
| "file": { |
| "class": "logging.FileHandler", |
| "filename": "runner.log", |
| "mode": "a", |
| "level": "DEBUG", |
| "formatter": "monai_default", |
| |
| }, |
| }, |
| } |
|
|
| if log_output_file is not None: |
| log_config["handlers"]["file"]["filename"] = log_output_file |
| log_config["handlers"]["file"]["level"] = "DEBUG" |
| else: |
| log_config["handlers"]["file"]["level"] = "CRITICAL" |
|
|
| if debug or bool(os.environ.get("SEGRESNET_DEBUG", False)): |
| log_config["handlers"]["console"]["level"] = "DEBUG" |
|
|
| logging.config.dictConfig(log_config) |
| |
| |
|
|
|
|
| def get_gpu_mem_size(): |
| gpu_mem = 0 |
| n_gpus = torch.cuda.device_count() |
| if n_gpus > 0: |
| gpu_mem = min([torch.cuda.get_device_properties(i).total_memory for i in range(n_gpus)]) |
| gpu_mem = gpu_mem / 1024**3 |
| else: |
| gpu_mem = 16 |
|
|
| return gpu_mem |
|
|
|
|
| def auto_adjust_network_settings( |
| auto_scale_roi=False, |
| auto_scale_batch=False, |
| auto_scale_filters=False, |
| image_size_mm=None, |
| spacing=None, |
| output_classes=None, |
| levels=None, |
| anisotropic_scales=False, |
| levels_limit=5, |
| gpu_mem=None, |
| ): |
| global_rank = 0 |
| if dist.is_available() and dist.is_initialized(): |
| global_rank = dist.get_rank() |
| print(f"auto_adjust_network_settings dist global_rank {global_rank}") |
| else: |
| print(f"auto_adjust_network_settings no distributed global_rank {global_rank}") |
|
|
| batch_size_default = 1 |
| init_filters_default = 32 |
|
|
| roi_size = np.array(roi_size_default) |
| base_numel = roi_size.prod() |
| gpu_factor = 1 |
|
|
| if gpu_mem is None: |
| gpu_mem = get_gpu_mem_size() |
| if global_rank == 0: |
| print(f"GPU device memory min: {gpu_mem}") |
|
|
| |
| if auto_scale_batch or auto_scale_roi or auto_scale_filters: |
| gpu_factor_init = gpu_factor = max(1, gpu_mem / 16) |
| if anisotropic_scales: |
| gpu_factor = max(1, 0.8 * gpu_factor) |
| if global_rank == 0: |
| print(f"base_numel {base_numel} gpu_factor {gpu_factor} gpu_factor_init {gpu_factor_init}") |
| else: |
| gpu_mem = 16 |
| gpu_factor = gpu_factor_init = 1 |
|
|
| |
| output_classes_thresh = 20 |
| if output_classes is not None and output_classes > output_classes_thresh: |
| base_adjust = gpu_mem / (output_classes * 0.2 + 11.5) |
| if gpu_mem < 17: |
| base_adjust /= 2 |
|
|
| if global_rank == 0: |
| print(f"base_adjust {base_adjust} since output_classes {output_classes} > {output_classes_thresh}") |
| if base_adjust < 0.95: |
| base_numel *= base_adjust |
| r = int(base_numel ** (1 / 3) / 2**4) |
| if r == 0 and global_rank == 0: |
| print(f"Warning: given output_classes {output_classes}, unable to fit any ROI on the gpu {gpu_mem} Gb!") |
| roi_size = np.array([max(1, r) * 2**4] * 3) |
| gpu_factor = gpu_factor_init = 1 |
| auto_scale_roi = False |
| else: |
| gpu_factor_init = gpu_factor = base_adjust |
|
|
| if global_rank == 0: |
| print(f"base_numel {base_numel} roi_size {roi_size} gpu_factor {gpu_factor}") |
|
|
| if image_size_mm is not None and spacing is not None: |
| image_size = np.floor(np.array(image_size_mm) / np.array(spacing)) |
| if global_rank == 0: |
| print(f"input roi {roi_size} image_size {image_size} numel {roi_size.prod()}") |
| roi_size = np.minimum(roi_size, image_size) |
| else: |
| raise ValueError("image_size_mm or spacing is not provided, network params may be inaccuracy") |
|
|
| |
| max_numel = base_numel * gpu_factor if auto_scale_roi else base_numel |
| while roi_size.prod() < max_numel: |
| old_numel = roi_size.prod() |
| roi_size = np.minimum(roi_size * 1.15, image_size) |
| if global_rank == 0: |
| print(f"increasing roi step {roi_size}") |
| if roi_size.prod() == old_numel: |
| break |
| if global_rank == 0: |
| print(f"increasing roi result 1 {roi_size}") |
|
|
| |
| if not anisotropic_scales: |
| if levels is None: |
| levels = np.floor(np.log2(roi_size)) |
| if global_rank == 0: |
| print(f"levels 1 {levels}") |
| levels = min(min(levels), levels_limit) |
| if global_rank == 0: |
| print(f"levels 2' {levels}") |
|
|
| factor = 2 ** (levels - 1) |
| roi_size = factor * np.maximum(2, np.floor(roi_size / factor)) |
| if global_rank == 0: |
| print(f"roi_size factored {roi_size}") |
|
|
| else: |
| extra_levels = np.floor(np.log2(np.max(spacing) / spacing)) |
| extra_levels = max(extra_levels) - extra_levels |
|
|
| if levels is None: |
| |
| levels = np.floor(np.log2(roi_size)) |
| if global_rank == 0: |
| print(f"levels 1 aniso {levels} extra_levels {extra_levels}") |
| levels = min(min(levels + extra_levels), levels_limit) |
| if global_rank == 0: |
| print(f"levels 2 {levels}") |
|
|
| factor = 2 ** (np.maximum(1, levels - extra_levels) - 1) |
| roi_size = factor * np.maximum(2, np.floor(roi_size / factor)) |
| if global_rank == 0: |
| print(f"roi_size factored {roi_size} factor {factor} extra_levels {extra_levels}") |
|
|
| |
| if auto_scale_filters and roi_size.prod() < base_numel * gpu_factor: |
| init_filters = int(max(32, np.floor(4 * (base_numel / roi_size.prod())) * 8)) |
| if global_rank == 0: |
| print(f"checking to increase init_filters {init_filters}") |
| gpu_factor_init *= init_filters / 32 |
| gpu_factor *= init_filters / 32 |
| else: |
| if global_rank == 0: |
| print(f"kept filters the same base_numel {base_numel}, gpu_factor {gpu_factor}") |
|
|
| init_filters = init_filters_default |
|
|
| |
| if auto_scale_batch and roi_size.prod() < base_numel * gpu_factor_init: |
| batch_size = int(1.1 * gpu_factor_init) |
| if global_rank == 0: |
| print( |
| f"increased batch_size {batch_size} base_numel {base_numel}, gpu_factor {gpu_factor}, gpu_factor_init {gpu_factor_init}" |
| ) |
|
|
| else: |
| batch_size = batch_size_default |
| if global_rank == 0: |
| print( |
| f"kept batch the same base_numel {base_numel}, gpu_factor {gpu_factor}, gpu_factor_init {gpu_factor_init}" |
| ) |
|
|
| levels = int(levels) |
| roi_size = roi_size.astype(int).tolist() |
|
|
| if global_rank == 0: |
| print( |
| f"Suggested network parameters: \n" |
| f"Batch size {batch_size_default} => {batch_size} \n" |
| f"ROI size {roi_size_default} => {roi_size} \n" |
| f"init_filters {init_filters_default} => {init_filters} \n" |
| f"aniso: {anisotropic_scales} image_size_mm: {image_size_mm} spacing: {spacing} levels: {levels} \n" |
| ) |
|
|
| return roi_size, levels, init_filters, batch_size |
|
|
