model-constructor 0.0.1

Constructor for pytorch models.

model_constructor

Constructor to create pytorch model.

_

Install

pip install model_constructor

How to use

model = Net()

model = Net()

model

Net(
  (stem): Stem(
    sizes: [3, 64]
    (conv0): ConvLayer(
      (conv): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
      (act_fn): ReLU(inplace=True)
      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  )
  (body): Body(
    (layer_0): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_1): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_2): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_3): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
  )
  (head): Head(
    (pool): AdaptiveAvgPool2d(output_size=(1, 1))
    (flat): Flatten()
    (fc): Linear(in_features=512, out_features=1000, bias=True)
  )
)

Resnet as example

Predefined Resnet18 and Resnet34

model = resnet34(num_classes=10)

model

Net(
  (stem): Stem(
    sizes: [3, 64]
    (conv0): ConvLayer(
      (conv): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
      (act_fn): ReLU(inplace=True)
      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  )
  (body): Body(
    (layer_0): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_2): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_1): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_2): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_3): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_2): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_2): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_3): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_4): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_5): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_3): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_2): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
  )
  (head): Head(
    (pool): AdaptiveAvgPool2d(output_size=(1, 1))
    (flat): Flatten()
    (fc): Linear(in_features=512, out_features=10, bias=True)
  )
)

But it can be created as:

If you need, you can customize it, for example

my_resnet18 = Net(block=BasicBlock, blocks=[2, 2, 2, 2])

my_resnet18

Net(
  (stem): Stem(
    sizes: [3, 64]
    (conv0): ConvLayer(
      (conv): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
      (act_fn): ReLU(inplace=True)
      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  )
  (body): Body(
    (layer_0): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_1): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_2): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_3): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
  )
  (head): Head(
    (pool): AdaptiveAvgPool2d(output_size=(1, 1))
    (flat): Flatten()
    (fc): Linear(in_features=512, out_features=1000, bias=True)
  )
)

my_resnet34 = Net(block=BasicBlock, blocks=[3, 4, 6, 3])

my_resnet34

Net(
  (stem): Stem(
    sizes: [3, 64]
    (conv0): ConvLayer(
      (conv): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
      (act_fn): ReLU(inplace=True)
      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  )
  (body): Body(
    (layer_0): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_2): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_1): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_2): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_3): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_2): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_2): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_3): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_4): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_5): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_3): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_2): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
  )
  (head): Head(
    (pool): AdaptiveAvgPool2d(output_size=(1, 1))
    (flat): Flatten()
    (fc): Linear(in_features=512, out_features=1000, bias=True)
  )
)

Some examples

Wee can experiment with models by changing some parts of model. Here only base functiomnality, but it can be easily extanded.

Here some examples:

Custom stem

model = Net(stem=partial(Stem, sizes=[32, 32]))

model.stem

Stem(
  sizes: [3, 32, 32, 64]
  (conv0): ConvLayer(
    (conv): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
    (act_fn): ReLU(inplace=True)
    (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  )
  (conv1): ConvLayer(
    (conv): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (act_fn): ReLU(inplace=True)
    (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  )
  (conv2): ConvLayer(
    (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (act_fn): ReLU(inplace=True)
    (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  )
  (pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
)

Act fn before Normalization

model = Net(bn_1st=False)

model

Net(
  (stem): Stem(
    sizes: [3, 64]
    (conv0): ConvLayer(
      (conv): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
      (act_fn): ReLU(inplace=True)
      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  )
  (body): Body(
    (layer_0): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_1): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_2): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_3): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): ReLU(inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
  )
  (head): Head(
    (pool): AdaptiveAvgPool2d(output_size=(1, 1))
    (flat): Flatten()
    (fc): Linear(in_features=512, out_features=1000, bias=True)
  )
)

Change activation function

act_fn = nn.LeakyReLU(inplace=True)

model = Net(act_fn=act_fn)

model

Net(
  (stem): Stem(
    sizes: [3, 64]
    (conv0): ConvLayer(
      (conv): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
      (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (pool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  )
  (body): Body(
    (layer_0): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_1): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_2): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
    (layer_3): Sequential(
      (block_0): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (downsample): ConvLayer(
          (conv): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2))
          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
      (block_1): BasicBlock(
        (conv): Sequential(
          (conv_0): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
          (conv_1): ConvLayer(
            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            (act_fn): LeakyReLU(negative_slope=0.01, inplace=True)
            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          )
        )
        (merge): Noop()
        (act_conn): ReLU(inplace=True)
      )
    )
  )
  (head): Head(
    (pool): AdaptiveAvgPool2d(output_size=(1, 1))
    (flat): Flatten()
    (fc): Linear(in_features=512, out_features=1000, bias=True)
  )
)