code: def compute_approx_vram_consumption(
      patch_size,
      num_pool_per_axis,
      base_num_features,
      max_num_features,
      num_modalities,
      num_classes,
      pool_op_kernel_sizes,
      deep_supervision=False,
      conv_per_stage=2,
  ):
      if not isinstance(num_pool_per_axis, np.ndarray):
          num_pool_per_axis = np.array(num_pool_per_axis)
  
      npool = len(pool_op_kernel_sizes)
  
      map_size = np.array(patch_size)
      tmp = np.int64(
          (conv_per_stage * 2 + 1)
          * np.prod(map_size, dtype=np.int64)
          * base_num_features
          + num_modalities * np.prod(map_size, dtype=np.int64)
          + num_classes * np.prod(map_size, dtype=np.int64)
      )
  
      num_feat = base_num_features
  
      for p in range(npool):
          for pi in range(len(num_pool_per_axis)):
              map_size[pi] /= pool_op_kernel_sizes[p][pi]
          num_feat = min(num_feat * 2, max_num_features)
          num_blocks = (
              (conv_per_stage * 2 + 1)
              if p < (npool - 1)
              else conv_per_stage
          )  # conv_per_stage + conv_per_stage for the convs of encode/decode and 1 for transposed conv
          tmp += (
              num_blocks * np.prod(map_size, dtype=np.int64) * num_feat
          )
          if deep_supervision and p < (npool - 2):
              tmp += np.prod(map_size, dtype=np.int64) * num_classes
      return tmp
  
file_path: PaddleSeg/contrib/MedicalSeg/tools/preprocess_utils/experiment_utils.py
imports: import numpy
  from copy import deepcopy
label_desc: Use image node_id 6 for calling a global function with example usage: compute_approx_vram_consumption(patch_size, num_pool_per_axis, base_num_features, max_num_features, num_modalities, num_classes, pool_op_kernel_sizes, deep_supervision, conv_per_stage) and returns: tmp