class_docstr: Env for channel pruning search.
  This class is used to prune model using specified pruner. It prunes one layer when
  step() is called. When the last layer is pruned, it evaluate the pruned model using
  evaluator, and use the returned value of evaluator as reward of the episode.
  
  Usage:
      env = ChannelPruningEnv(pruner, evaluator, val_loader, checkpoint, env_args)
      episode = 0
      T = []
      while episode < num_episode:
          action = agent.select_action(observation)
          observation2, reward, done, info = env.step(action)
          T.append([reward, deepcopy(observation), deepcopy(observation2), action, done])
  
          if done: # end of episode, last layer pruned
              episode += 1
              # train agent with episode data
              for _, s_t, s_t1, a_t, done in T:
                  agent.observe(final_reward, s_t, s_t1, a_t, done)
                  agent.update_policy()
              T = []
  
  Attributes:
      prunable_idx: layer indices for pruable layers, the index values are the index
          of list(self.model.modules()). Pruable layers are pointwise Conv2d layers and Linear
          layers.
      buffer_idx: layer indices for buffer layers which refers the depthwise layers.
          Each depthwise layer is always followd by a pointwise layer for both mobilenet and
          mobilenetv2. The depthwise layer's filters are pruned when its next pointwise layer's
          corresponding input channels are pruned.
      shared_idx: layer indices for layers which share input.
          For example: [[1,4], [8, 10, 15]] means layer 1 and 4 share same input, and layer
          8, 10 and 15 share another input.
      layer_embedding: embeddings for each prunable layers, the embedding is used as
          observation for DDPG agent.
      layer_info_dict: flops and number of parameters of each layer.
      min_strategy_dict: key is layer index, value is a tuple, the first value is the minimum
          action of input channel, the second value is the minimum action value of output channel.
      strategy_dict: key is layer index, value is a tuple, the first value is the action of input
          channel, the second value is the action of output channel.
  
  Parameters:
      pruner: Pruner
          NNI Pruner instance used to prune model.
      evaluator: function
          function to evaluate the pruned model.
          The prototype of the function:
              >>> def evaluator(val_loader, model):
              >>>     ...
              >>>     return acc
      val_loader: torch.utils.data.DataLoader
          Data loader of validation dataset.
      checkpoint: dict
          checkpoint of the model to be pruned. It is used to reset model at beginning of each
          episode.
      args:
          A Namespace object containing following arguments:
          model_type: str
              model type to prune, currently 'mobilenet' and 'mobilenetv2' are supported.
          flops_ratio: float
              preserve flops ratio.
          lbound: float
              minimum weight preserve ratio for each layer.
          rbound: float
              maximum weight preserve ratio for each layer.
          reward: function
              reward function type
  
          # parameters for channel pruning
          n_calibration_batches: int
              number of batches to extract layer information.
          n_points_per_layer: int
              number of feature points per layer.
          channel_round: int
              round channel to multiple of channel_round.
code: def _extract_layer_information(self):
      m_list = list(self.model.modules())
  
      self.data_saver = []
      self.layer_info_dict = dict()
      self.wsize_list = []
      self.flops_list = []
  
      from .lib.utils import measure_layer_for_pruning
  
      # extend the forward fn to record layer info
      def new_forward(m):
          def lambda_forward(x):
              m.input_feat = x.clone()
              # TODO replace this flops counter with nni.compression.torch.utils.counter.count_flops_params
              measure_layer_for_pruning(m, x)
              y = m.old_forward(x)
              m.output_feat = y.clone()
              return y
  
          return lambda_forward
  
      device = None
      for idx in self.prunable_idx + self.buffer_idx:  # get all
          m = m_list[idx]
          m.old_forward = m.forward
          m.forward = new_forward(m)
          if device is None and type(m) == PrunerModuleWrapper:
              device = m.module.weight.device
  
      # now let the image flow
      _logger.info("=> Extracting information...")
      with torch.no_grad():
          for i_b, (inputs, target) in enumerate(
              self._val_loader
          ):  # use image from train set
              if i_b == self.n_calibration_batches:
                  break
              self.data_saver.append((inputs.clone(), target.clone()))
              input_var = torch.autograd.Variable(inputs).to(device)
  
              # inference and collect stats
              _ = self.model(input_var)
  
              if i_b == 0:  # first batch
                  for idx in self.prunable_idx + self.buffer_idx:
                      self.layer_info_dict[idx] = dict()
                      self.layer_info_dict[idx]["params"] = m_list[
                          idx
                      ].params
                      self.layer_info_dict[idx]["flops"] = m_list[
                          idx
                      ].flops
                      self.wsize_list.append(m_list[idx].params)
                      self.flops_list.append(m_list[idx].flops)
                  _logger.info("flops: %s", self.flops_list)
              for idx in self.prunable_idx:
                  f_in_np = m_list[idx].input_feat.data.cpu().numpy()
                  f_out_np = m_list[idx].output_feat.data.cpu().numpy()
                  if len(f_in_np.shape) == 4:  # conv
                      if (
                          self.prunable_idx.index(idx) == 0
                      ):  # first conv
                          f_in2save, f_out2save = None, None
                      elif (
                          m_list[idx].module.weight.size(3) > 1
                      ):  # normal conv
                          f_in2save, f_out2save = f_in_np, f_out_np
                      else:  # 1x1 conv
                          # assert f_out_np.shape[2] == f_in_np.shape[2]  # now support k=3
                          randx = np.random.randint(
                              0,
                              f_out_np.shape[2] - 0,
                              self.n_points_per_layer,
                          )
                          randy = np.random.randint(
                              0,
                              f_out_np.shape[3] - 0,
                              self.n_points_per_layer,
                          )
                          # input: [N, C, H, W]
                          self.layer_info_dict[idx][
                              (i_b, "randx")
                          ] = randx.copy()
                          self.layer_info_dict[idx][
                              (i_b, "randy")
                          ] = randy.copy()
  
                          f_in2save = (
                              f_in_np[:, :, randx, randy]
                              .copy()
                              .transpose(0, 2, 1)
                              .reshape(
                                  self.batch_size
                                  * self.n_points_per_layer,
                                  -1,
                              )
                          )
  
                          f_out2save = (
                              f_out_np[:, :, randx, randy]
                              .copy()
                              .transpose(0, 2, 1)
                              .reshape(
                                  self.batch_size
                                  * self.n_points_per_layer,
                                  -1,
                              )
                          )
                  else:
                      assert len(f_in_np.shape) == 2
                      f_in2save = f_in_np.copy()
                      f_out2save = f_out_np.copy()
                  if "input_feat" not in self.layer_info_dict[idx]:
                      self.layer_info_dict[idx][
                          "input_feat"
                      ] = f_in2save
                      self.layer_info_dict[idx][
                          "output_feat"
                      ] = f_out2save
                  else:
                      self.layer_info_dict[idx][
                          "input_feat"
                      ] = np.vstack(
                          (
                              self.layer_info_dict[idx]["input_feat"],
                              f_in2save,
                          )
                      )
                      self.layer_info_dict[idx][
                          "output_feat"
                      ] = np.vstack(
                          (
                              self.layer_info_dict[idx]["output_feat"],
                              f_out2save,
                          )
                      )
  
imports: import os
  import logging
  import time
  import math
  import copy
  import numpy
  import torch
  import torch.nn
  from ...compressor import PrunerModuleWrapper
  from ..None import AMCWeightMasker
label_desc: Use image node_id 11 for calling the ChannelPruningEnv obj's underlying member method code with example usage: obj._extract_layer_information() and returns: lambda_forward, y