class_docstr: Convolutional neural network (CNN) model in pytorch.
  
  Args:
      spec (garage.InOutSpec): Specification of inputs and outputs.
          The input should be in 'NCHW' format: [batch_size, channel, height,
          width]. Will print a warning if the channel size is not 1 or 3.
          If output_space is specified, then a final linear layer will be
          inserted to map to that dimensionality.
          If output_space is None, it will be filled in with the computed
          output space.
      image_format (str): Either 'NCHW' or 'NHWC'. Should match the input
          specification. Gym uses NHWC by default, but PyTorch uses NCHW by
          default.
      hidden_channels (tuple[int]): Number of output channels for CNN.
          For example, (3, 32) means there are two convolutional layers.
          The filter for the first conv layer outputs 3 channels
          and the second one outputs 32 channels.
      kernel_sizes (tuple[int]): Dimension of the conv filters.
          For example, (3, 5) means there are two convolutional layers.
          The filter for first layer is of dimension (3 x 3)
          and the second one is of dimension (5 x 5).
      strides (tuple[int]): The stride of the sliding window. For example,
          (1, 2) means there are two convolutional layers. The stride of the
          filter for first layer is 1 and that of the second layer is 2.
      paddings (tuple[int]): Amount of zero-padding added to both sides of
          the input of a conv layer.
      padding_mode (str): The type of padding algorithm to use, i.e.
          'constant', 'reflect', 'replicate' or 'circular' and
          by default is 'zeros'.
      hidden_nonlinearity (callable or torch.nn.Module):
          Activation function for intermediate dense layer(s).
          It should return a torch.Tensor. Set it to None to maintain a
          linear activation.
      hidden_b_init (callable): Initializer function for the bias
          of intermediate dense layer(s). The function should return a
          torch.Tensor.
      max_pool (bool): Bool for using max-pooling or not.
      pool_shape (tuple[int]): Dimension of the pooling layer(s). For
          example, (2, 2) means that all pooling layers are of the same
          shape (2, 2).
      pool_stride (tuple[int]): The strides of the pooling layer(s). For
          example, (2, 2) means that all the pooling layers have
          strides (2, 2).
      layer_normalization (bool): Bool for using layer normalization or not.
      hidden_w_init (callable): Initializer function for the weight
          of intermediate dense layer(s). The function should return a
          torch.Tensor.
      enable_cudnn_benchmarks (bool): Whether to enable cudnn benchmarks
          in `torch`. If enabled, the backend selects the CNN benchamark
          algorithm with the best performance.
class_docstr_tok: ["Convolutional","neural","network","(","CNN",")","model","in","pytorch",".","Args",":","spec","(","garage.InOutSpec",")",":","Specification","of","inputs","and","outputs",".","The","input","should","be","in","'NCHW","'","format",":","[","batch_size",",","channel",",","height",",","width","]",".","Will","print","a","warning","if","the","channel","size","is","not","1","or","3",".","If","output_space","is","specified",",","then","a","final","linear","layer","will","be","inserted","to","map","to","that","dimensionality",".","If","output_space","is","None",",","it","will","be","filled","in","with","the","computed","output","space",".","image_format","(","str",")",":","Either","'NCHW","'","or","'NHWC","'",".","Should","match","the","input","specification",".","Gym","uses","NHWC","by","default",",","but","PyTorch","uses","NCHW","by","default",".","hidden_channels","(","tuple","[","int","]",")",":","Number","of","output","channels","for","CNN",".","For","example",",","(","3",",","32",")","means","there","are","two","convolutional","layers",".","The","filter","for","the","first","conv","layer","outputs","3","channels","and","the","second","one","outputs","32","channels",".","kernel_sizes","(","tuple","[","int","]",")",":","Dimension","of","the","conv","filters",".","For","example",",","(","3",",","5",")","means","there","are","two","convolutional","layers",".","The","filter","for","first","layer","is","of","dimension","(","3","x","3",")","and","the","second","one","is","of","dimension","(","5","x","5",")",".","strides","(","tuple","[","int","]",")",":","The","stride","of","the","sliding","window",".","For","example",",","(","1",",","2",")","means","there","are","two","convolutional","layers",".","The","stride","of","the","filter","for","first","layer","is","1","and","that","of","the","second","layer","is","2.","paddings","(","tuple","[","int","]",")",":","Amount","of","zero-padding","added","to","both","sides","of","the","input","of","a","conv","layer",".","padding_mode","(","str",")",":","The","type","of","padding","algorithm","to","use",",","i.e",".","'constant","'",",","'reflect","'",",","'replicate","'","or","'circular","'","and","by","default","is","'zeros","'",".","hidden_nonlinearity","(","callable","or","torch.nn.Module",")",":","Activation","function","for","intermediate","dense","layer","(","s",")",".","It","should","return","a","torch.Tensor",".","Set","it","to","None","to","maintain","a","linear","activation",".","hidden_b_init","(","callable",")",":","Initializer","function","for","the","bias","of","intermediate","dense","layer","(","s",")",".","The","function","should","return","a","torch.Tensor",".","max_pool","(","bool",")",":","Bool","for","using","max-pooling","or","not",".","pool_shape","(","tuple","[","int","]",")",":","Dimension","of","the","pooling","layer","(","s",")",".","For","example",",","(","2",",","2",")","means","that","all","pooling","layers","are","of","the","same","shape","(","2",",","2",")",".","pool_stride","(","tuple","[","int","]",")",":","The","strides","of","the","pooling","layer","(","s",")",".","For","example",",","(","2",",","2",")","means","that","all","the","pooling","layers","have","strides","(","2",",","2",")",".","layer_normalization","(","bool",")",":","Bool","for","using","layer","normalization","or","not",".","hidden_w_init","(","callable",")",":","Initializer","function","for","the","weight","of","intermediate","dense","layer","(","s",")",".","The","function","should","return","a","torch.Tensor",".","enable_cudnn_benchmarks","(","bool",")",":","Whether","to","enable","cudnn","benchmarks","in","`","torch","`",".","If","enabled",",","the","backend","selects","the","CNN","benchamark","algorithm","with","the","best","performance","."]
code: def __init__(
      self,
      spec,
      image_format,
      hidden_channels,
      *,  # Many things after this are ints or tuples of ints.
      kernel_sizes,
      strides,
      paddings=0,
      padding_mode="zeros",
      hidden_nonlinearity=nn.ReLU,
      hidden_w_init=nn.init.xavier_uniform_,
      hidden_b_init=nn.init.zeros_,
      max_pool=False,
      pool_shape=None,
      pool_stride=1,
      layer_normalization=False,
      enable_cudnn_benchmarks=True,
  ):
      super().__init__()
      assert len(hidden_channels) > 0
      # PyTorch forces us to use NCHW internally.
      in_channels, height, width = _check_spec(spec, image_format)
      self._format = image_format
      kernel_sizes = expand_var(
          "kernel_sizes",
          kernel_sizes,
          len(hidden_channels),
          "hidden_channels",
      )
      strides = expand_var(
          "strides", strides, len(hidden_channels), "hidden_channels"
      )
      paddings = expand_var(
          "paddings", paddings, len(hidden_channels), "hidden_channels"
      )
      pool_shape = expand_var(
          "pool_shape",
          pool_shape,
          len(hidden_channels),
          "hidden_channels",
      )
      pool_stride = expand_var(
          "pool_stride",
          pool_stride,
          len(hidden_channels),
          "hidden_channels",
      )
  
      self._cnn_layers = nn.Sequential()
      torch.backends.cudnn.benchmark = enable_cudnn_benchmarks
  
      # In case there are no hidden channels, handle output case.
      out_channels = in_channels
      for i, out_channels in enumerate(hidden_channels):
          conv_layer = nn.Conv2d(
              in_channels=in_channels,
              out_channels=out_channels,
              kernel_size=kernel_sizes[i],
              stride=strides[i],
              padding=paddings[i],
              padding_mode=padding_mode,
          )
          height = output_height_2d(conv_layer, height)
          width = output_width_2d(conv_layer, width)
          hidden_w_init(conv_layer.weight)
          hidden_b_init(conv_layer.bias)
          self._cnn_layers.add_module(f"conv_{i}", conv_layer)
  
          if layer_normalization:
              self._cnn_layers.add_module(
                  f"layer_norm_{i}",
                  nn.LayerNorm((out_channels, height, width)),
              )
  
          if hidden_nonlinearity:
              self._cnn_layers.add_module(
                  f"non_linearity_{i}",
                  NonLinearity(hidden_nonlinearity),
              )
  
          if max_pool:
              pool = nn.MaxPool2d(
                  kernel_size=pool_shape[i], stride=pool_stride[i]
              )
              height = output_height_2d(pool, height)
              width = output_width_2d(pool, width)
              self._cnn_layers.add_module(f"max_pooling_{i}", pool)
  
          in_channels = out_channels
  
      output_dims = out_channels * height * width
  
      if spec.output_space is None:
          final_spec = InOutSpec(
              spec.input_space,
              akro.Box(low=-np.inf, high=np.inf, shape=(output_dims,)),
          )
          self._final_layer = None
      else:
          final_spec = spec
          # Checked at start of __init__
          self._final_layer = nn.Linear(
              output_dims, spec.output_space.shape[0]
          )
  
      self.spec = final_spec
  
file_path: garage/src/garage/torch/modules/cnn_module.py
imports: import warnings
  import akro
  import numpy
  import torch
  from torch import nn
  from garage import InOutSpec
  from garage.torch import expand_var, NonLinearity, output_height_2d, output_width_2d
label_desc: Use image node_id 1 to create a new CNNModule object from inherited base classes: nn with example: obj = CNNModule(spec, image_format, hidden_channels)