class_docstr: A StrategyGenerator which deals with the sharding strategies of batch normalization.
  
  To keep the math consistency, there are two way to do BatchNorm if the input
  shards on batch dimension:
  1. We gather the input partitions through batch dimension, then do the normal BatchNorm.
  2. We do the SyncBatchNorm on the each input partition separately, the SyncBN op will help
     us to keep the computing correctness.
  In this generator, both methods will be considered.
class_docstr_tok: ["A","StrategyGenerator","which","deals","with","the","sharding","strategies","of","batch","normalization",".","To","keep","the","math","consistency",",","there","are","two","way","to","do","BatchNorm","if","the","input","shards","on","batch","dimension",":","1",".","We","gather","the","input","partitions","through","batch","dimension",",","then","do","the","normal","BatchNorm",".","2",".","We","do","the","SyncBatchNorm","on","the","each","input","partition","separately",",","the","SyncBN","op","will","help","us","to","keep","the","computing","correctness",".","In","this","generator",",","both","methods","will","be","considered","."]
code: def update_memory_cost(self, strategy: ShardingStrategy):
      forward_size_mapping = {
          "input": self._compute_size_in_bytes(strategy, "input"),
          "other": self._compute_size_in_bytes(strategy, "other"),
          "output": self._compute_size_in_bytes(strategy, "output"),
          "running_mean": self._compute_size_in_bytes(
              strategy, "running_mean"
          ),
          "running_var": self._compute_size_in_bytes(
              strategy, "running_var"
          ),
      }
  
      if self.has_bias:
          bias_size = self._compute_size_in_bytes(strategy, "bias")
          forward_size_mapping["bias"] = bias_size
  
      backward_size_mapping = copy.deepcopy(forward_size_mapping)
      backward_size_mapping.pop("output")
      # compute fwd cost incurred
      # fwd_cost = input + other + bias + output
      fwd_activation_cost = sum(
          [
              v
              for k, v in forward_size_mapping.items()
              if not self.is_param(k) and not self.is_buffer(k)
          ]
      )
      fwd_parameter_cost = sum(
          [
              v
              for k, v in forward_size_mapping.items()
              if self.is_param(k)
          ]
      )
      fwd_buffer_cost = sum(
          [
              v
              for k, v in forward_size_mapping.items()
              if self.is_buffer(k)
          ]
      )
      fwd_mem_cost = MemoryCost(
          activation=fwd_activation_cost,
          parameter=fwd_parameter_cost,
          buffer=fwd_buffer_cost,
      )
  
      # compute bwd cost incurred
      # bwd_cost = input_grad + other_grad + bias_grad
      bwd_activation_cost = sum(
          [
              v
              for k, v in backward_size_mapping.items()
              if not self.is_param(k) and not self.is_buffer(k)
          ]
      )
      bwd_parameter_cost = sum(
          [
              v
              for k, v in backward_size_mapping.items()
              if self.is_param(k)
          ]
      )
      bwd_mem_cost = MemoryCost(
          activation=bwd_activation_cost, parameter=bwd_parameter_cost
      )
  
      # compute total cost
      total_mem_cost = MemoryCost(
          activation=fwd_activation_cost + bwd_activation_cost,
          parameter=fwd_parameter_cost + bwd_parameter_cost,
          buffer=fwd_buffer_cost,
      )
      memory_cost = TrainCycleItem(
          fwd=fwd_mem_cost, bwd=bwd_mem_cost, total=total_mem_cost
      )
      strategy.memory_cost = memory_cost
  
imports: import copy
  import operator
  from functools import reduce
  from typing import List
  from colossalai.auto_parallel.tensor_shard.sharding_strategy import CommType, MemoryCost, ShardingStrategy, TrainCycleItem
  from colossalai.auto_parallel.tensor_shard.utils import ignore_sharding_exception
  from colossalai.tensor.shape_consistency import CollectiveCommPattern
  from .strategy_generator import StrategyGenerator
label_desc: Use image node_id 3 for calling the BatchNormStrategyGenerator obj's underlying member method code with example usage: obj.update_memory_cost(strategy) without return types