class_docstr: pytorch 分类 PbaAugment transform
class_docstr_tok: ["pytorch","\u5206\u7c7b","PbaAugment","transform"]
docstr: Args:
      input_size:
      scale_size:
      normalize:
      pre_transform:
      **kwargs:
docstr_tok: ["Args",":","input_size",":","scale_size",":","normalize",":","pre_transform",":","*","*","kwargs",":"]
code: def __init__(
      self,
      input_size: int = 224,
      scale_size: int = 256,
      normalize: Optional[list] = None,
      pre_transform: bool = True,
      stage: str = "search",
      **kwargs,
  ) -> None:
      """
  
      Args:
          input_size:
          scale_size:
          normalize:
          pre_transform:
          **kwargs:
      """
  
      if normalize is None:
          normalize = {
              "mean": [0.485, 0.456, 0.406],
              "std": [0.229, 0.224, 0.225],
          }
  
      policy = kwargs["policy"]
      assert stage in ["search", "train"]
      train_epochs = kwargs["hp_policy_epochs"]
      self.auto_aug_transform = AutoAugTransform.create(
          policy, stage=stage, train_epochs=train_epochs
      )
      # self.auto_aug_transform = PbtAutoAugmentClassiferTransform(conf)
      if pre_transform:
          self.pre_transform = transforms.Resize(input_size)
  
      self.post_transform = transforms.Compose(
          transforms=[
              transforms.Permute(),
              transforms.Normalize(**normalize, channel_first=True),
          ],
          channel_first=False,
      )
      self.cur_epoch = 0
  
imports: import re
  import numpy
  from typing import Dict, List, Optional, Union, Tuple
  import six
  import cv2
  import os
  import paddle
  import paddlehub.vision.transforms
  from PIL import ImageFile
  from auto_augment.autoaug.transform.autoaug_transform import AutoAugTransform
label_desc: Use image node_id 1 to create a new PbaAugment object from inherited base classes: object with example: obj = PbaAugment(input_size, scale_size, normalize, pre_transform, stage)