metacontroller-pytorch 0.0.31__tar.gz → 0.0.33__tar.gz

{metacontroller_pytorch-0.0.31 → metacontroller_pytorch-0.0.33}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metacontroller-pytorch
-Version: 0.0.31
+Version: 0.0.33
 Summary: Transformer Metacontroller
 Project-URL: Homepage, https://pypi.org/project/metacontroller/
 Project-URL: Repository, https://github.com/lucidrains/metacontroller
@@ -40,7 +40,7 @@ Requires-Dist: einops>=0.8.1
 Requires-Dist: einx>=0.3.0
 Requires-Dist: loguru
 Requires-Dist: memmap-replay-buffer>=0.0.23
-Requires-Dist: torch-einops-utils>=0.0.16
+Requires-Dist: torch-einops-utils>=0.0.19
 Requires-Dist: torch>=2.5
 Requires-Dist: vector-quantize-pytorch>=1.27.20
 Requires-Dist: x-evolution>=0.1.23

{metacontroller_pytorch-0.0.31 → metacontroller_pytorch-0.0.33}/metacontroller/metacontroller.py

@@ -26,7 +26,7 @@ from discrete_continuous_embed_readout import Embed, Readout, EmbedAndReadout
 
 from assoc_scan import AssocScan
 
-from torch_einops_utils import maybe, pad_at_dim, lens_to_mask
+from torch_einops_utils import maybe, pad_at_dim, lens_to_mask, masked_mean, align_dims_left, pad_right_ndim_to
 from torch_einops_utils.save_load import save_load
 
 # constants
@@ -66,6 +66,47 @@ MetaControllerOutput = namedtuple('MetaControllerOutput', (
     'switch_loss'
 ))
 
+def z_score(t, eps = 1e-8):
+    return (t - t.mean()) / (t.std() + eps)
+
+def policy_loss(
+    meta_controller,
+    state,
+    old_log_probs,
+    actions,
+    advantages,
+    mask,
+    episode_lens = None,
+    eps_clip = 0.2
+):
+    # get new log probs
+
+    action_dist = meta_controller.get_action_dist_for_internal_rl(state)
+    new_log_probs = meta_controller.log_prob(action_dist, actions)
+
+    # calculate ratio
+
+    ratio = (new_log_probs - old_log_probs).exp()
+
+    # align ratio and advantages
+
+    ratio, advantages = align_dims_left((ratio, advantages))
+
+    # ppo surrogate loss
+
+    surr1 = ratio * advantages
+    surr2 = ratio.clamp(1 - eps_clip, 1 + eps_clip) * advantages
+
+    losses = -torch.min(surr1, surr2)
+
+    # masking
+
+    if exists(episode_lens):
+        mask, episode_mask = align_dims_left((mask, lens_to_mask(episode_lens, losses.shape[1])))
+        mask = mask & episode_mask
+
+    return masked_mean(losses, mask)
+
 @save_load()
 class MetaController(Module):
     def __init__(
@@ -107,7 +148,6 @@ class MetaController(Module):
 
         self.switch_per_latent_dim = switch_per_latent_dim
 
-
         self.dim_latent = dim_latent
         self.switching_unit = GRU(dim_meta + dim_latent, dim_meta)
         self.to_switching_unit_beta = nn.Linear(dim_meta, dim_latent if switch_per_latent_dim else 1, bias = False)
@@ -147,6 +187,23 @@ class MetaController(Module):
             *self.action_proposer_mean_log_var.parameters()
         ]
 
+    def get_action_dist_for_internal_rl(
+        self,
+        residual_stream
+    ):
+        meta_embed = self.model_to_meta(residual_stream)
+
+        proposed_action_hidden, _ = self.action_proposer(meta_embed)
+
+        return self.action_proposer_mean_log_var(proposed_action_hidden)
+
+    def log_prob(
+        self,
+        action_dist,
+        sampled_latent_action
+    ):
+        return self.action_proposer_mean_log_var.log_prob(action_dist, sampled_latent_action)
+
     def forward(
         self,
         residual_stream,
@@ -276,6 +333,12 @@ class MetaController(Module):
 
 # main transformer, which is subsumed into the environment after behavioral cloning
 
+Hiddens = namedtuple('Hiddens', (
+    'lower_body',
+    'meta_controller',
+    'upper_body'
+))
+
 TransformerOutput = namedtuple('TransformerOutput', (
     'residual_stream_latent',
     'prev_hiddens'
@@ -417,9 +480,8 @@ class Transformer(Module):
         # maybe return behavior cloning loss
 
         if behavioral_cloning:
-            loss_mask = None
-            if exists(episode_lens):
-                loss_mask = lens_to_mask(episode_lens, state.shape[1])
+
+            loss_mask = maybe(lens_to_mask)(episode_lens, state.shape[1])
 
             state_dist_params = self.state_readout(attended)
             state_clone_loss = self.state_readout.calculate_loss(state_dist_params, target_state, mask = loss_mask)
@@ -441,4 +503,4 @@ class Transformer(Module):
         if return_one:
             return dist_params
 
-        return dist_params, TransformerOutput(residual_stream, (next_lower_hiddens, next_meta_hiddens, next_upper_hiddens))
+        return dist_params, TransformerOutput(residual_stream, Hiddens(next_lower_hiddens, next_meta_hiddens, next_upper_hiddens))

{metacontroller_pytorch-0.0.31 → metacontroller_pytorch-0.0.33}/metacontroller/metacontroller_with_binary_mapper.py

@@ -23,7 +23,7 @@ from x_mlps_pytorch import Feedforwards
 
 from assoc_scan import AssocScan
 
-from torch_einops_utils import maybe, pad_at_dim, lens_to_mask, align_dims_left
+from torch_einops_utils import maybe, pad_at_dim, lens_to_mask, masked_mean, align_dims_left, pad_right_ndim_to
 from torch_einops_utils.save_load import save_load
 
 from vector_quantize_pytorch import BinaryMapper
@@ -50,6 +50,9 @@ def default(*args):
 def straight_through(src, tgt):
     return tgt + src - src.detach()
 
+def log(t, eps = 1e-20):
+    return t.clamp_min(eps).log()
+
 # meta controller
 
 @save_load()
@@ -71,6 +74,9 @@ class MetaControllerWithBinaryMapper(Module):
         kl_loss_threshold = 0.
     ):
         super().__init__()
+
+        assert not switch_per_code, 'switch_per_code is not supported for binary mapper'
+
         dim_meta = default(dim_meta_controller, dim_model)
 
         self.model_to_meta = Linear(dim_model, dim_meta)
@@ -137,6 +143,23 @@ class MetaControllerWithBinaryMapper(Module):
             *self.proposer_to_binary_logits.parameters()
         ]
 
+    def log_prob(
+        self,
+        action_dist,
+        sampled_latent_action
+    ):
+        action_prob = action_dist.sigmoid()
+        probs = stack((action_prob, 1. - action_prob), dim = -1)
+        log_probs = log(probs)
+
+        indices = sampled_latent_action.argmax(dim = -1)
+        codes = self.binary_mapper.codes[indices].long()
+
+        codes = rearrange(codes, '... -> ... 1')
+        action_log_probs = log_probs.gather(-1, codes)
+
+        return rearrange(action_log_probs, '... 1 -> ...')
+
     def forward(
         self,
         residual_stream,

metacontroller_pytorch-0.0.33/metacontroller/transformer_with_resnet.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+import torch
+from torch import nn, Tensor
+from torch.nn import Module, ModuleList
+from einops import rearrange
+from einops.layers.torch import Rearrange
+
+from metacontroller.metacontroller import Transformer
+
+from torch_einops_utils import pack_with_inverse
+
+# resnet components
+
+def exists(v):
+    return v is not None
+
+class BasicBlock(Module):
+    expansion = 1
+
+    def __init__(
+        self,
+        dim,
+        dim_out,
+        stride = 1,
+        downsample: Module | None = None
+    ):
+        super().__init__()
+        self.conv1 = nn.Conv2d(dim, dim_out, 3, stride = stride, padding = 1, bias = False)
+        self.bn1 = nn.BatchNorm2d(dim_out)
+        self.relu = nn.ReLU(inplace = True)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, padding = 1, bias = False)
+        self.bn2 = nn.BatchNorm2d(dim_out)
+        self.downsample = downsample
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if exists(self.downsample):
+            identity = self.downsample(x)
+
+        out += identity
+        return self.relu(out)
+
+class Bottleneck(Module):
+    expansion = 4
+
+    def __init__(
+        self,
+        dim,
+        dim_out,
+        stride = 1,
+        downsample: Module | None = None
+    ):
+        super().__init__()
+        width = dim_out # simple resnet shortcut
+        self.conv1 = nn.Conv2d(dim, width, 1, bias = False)
+        self.bn1 = nn.BatchNorm2d(width)
+        self.conv2 = nn.Conv2d(width, width, 3, stride = stride, padding = 1, bias = False)
+        self.bn2 = nn.BatchNorm2d(width)
+        self.conv3 = nn.Conv2d(width, dim_out * self.expansion, 1, bias = False)
+        self.bn3 = nn.BatchNorm2d(dim_out * self.expansion)
+        self.relu = nn.ReLU(inplace = True)
+        self.downsample = downsample
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if exists(self.downsample):
+            identity = self.downsample(x)
+
+        out += identity
+        return self.relu(out)
+
+class ResNet(Module):
+    def __init__(
+        self,
+        block: type[BasicBlock | Bottleneck],
+        layers: list[int],
+        num_classes = 1000,
+        channels = 3
+    ):
+        super().__init__()
+        self.inplanes = 64
+
+        self.conv1 = nn.Conv2d(channels, 64, kernel_size = 7, stride = 2, padding = 3, bias = False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace = True)
+        self.maxpool = nn.MaxPool2d(kernel_size = 3, stride = 2, padding = 1)
+
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride = 2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride = 2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride = 2)
+
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.flatten = Rearrange('b c 1 1 -> b c')
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+    def _make_layer(
+        self,
+        block: type[BasicBlock | Bottleneck],
+        planes: int,
+        blocks: int,
+        stride: int = 1
+    ) -> nn.Sequential:
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion, 1, stride = stride, bias = False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = self.flatten(x)
+        x = self.fc(x)
+        return x
+
+# resnet factory
+
+def resnet18(num_classes: any = 1000):
+    return ResNet(BasicBlock, [2, 2, 2, 2], num_classes)
+
+def resnet34(num_classes: any = 1000):
+    return ResNet(BasicBlock, [3, 4, 6, 3], num_classes)
+
+def resnet50(num_classes: any = 1000):
+    return ResNet(Bottleneck, [3, 4, 6, 3], num_classes)
+
+# transformer with resnet
+
+class TransformerWithResnet(Transformer):
+    def __init__(
+        self,
+        *,
+        resnet_type = 'resnet18',
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        resnet_klass = resnet18
+        if resnet_type == 'resnet34':
+            resnet_klass = resnet34
+        elif resnet_type == 'resnet50':
+            resnet_klass = resnet50
+
+        self.resnet_dim = kwargs['state_embed_readout']['num_continuous']
+        self.visual_encoder = resnet_klass(num_classes = self.resnet_dim)
+
+    def visual_encode(self, x: Tensor) -> Tensor:
+        if x.shape[-1] == 3:
+            x = rearrange(x, '... h w c -> ... c h w')
+
+        x, inverse = pack_with_inverse(x, '* c h w')
+
+        h = self.visual_encoder(x)
+
+        return inverse(h, '* d')

{metacontroller_pytorch-0.0.31 → metacontroller_pytorch-0.0.33}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "metacontroller-pytorch"
-version = "0.0.31"
+version = "0.0.33"
 description = "Transformer Metacontroller"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -31,7 +31,7 @@ dependencies = [
     "loguru",
     "memmap-replay-buffer>=0.0.23",
     "torch>=2.5",
-    "torch-einops-utils>=0.0.16",
+    "torch-einops-utils>=0.0.19",
     "vector-quantize-pytorch>=1.27.20",
     "x-evolution>=0.1.23",
     "x-mlps-pytorch",

{metacontroller_pytorch-0.0.31 → metacontroller_pytorch-0.0.33}/tests/test_metacontroller.py

@@ -4,19 +4,30 @@ param = pytest.mark.parametrize
 from pathlib import Path
 
 import torch
-from metacontroller.metacontroller import Transformer, MetaController
+from torch import cat
+from metacontroller.metacontroller import Transformer, MetaController, policy_loss, z_score
 from metacontroller.metacontroller_with_binary_mapper import MetaControllerWithBinaryMapper
 
 from einops import rearrange
 
-@param('use_binary_mapper_variant', (False, True))
+# functions
+
+def exists(v):
+    return v is not None
+
+# test
+
+@param('use_binary_mapper_variant, switch_per_latent_dim', [
+    (False, False),
+    (False, True),
+    (True, False)
+])
 @param('action_discrete', (False, True))
-@param('switch_per_latent_dim', (False, True))
 @param('variable_length', (False, True))
 def test_metacontroller(
     use_binary_mapper_variant,
-    action_discrete,
     switch_per_latent_dim,
+    action_discrete,
     variable_length
 ):
 
@@ -69,16 +80,77 @@ def test_metacontroller(
 
     # internal rl - done iteratively
 
-    cache = None
-    past_action_id = None
+    # simulate grpo
+
+    all_episodes = []
+    all_rewards = []
+
+    for _ in range(3): # group of 3
+        subset_state = state[:1]
+
+        cache = None
+        past_action_id = None
+
+        states = []
+        log_probs = []
+        switch_betas = []
+        latent_actions = []
+
+        for one_state in subset_state.unbind(dim = 1):
+            one_state = rearrange(one_state, 'b d -> b 1 d')
+
+            logits, cache = model(one_state, past_action_id, meta_controller = meta_controller, return_cache = True)
+
+            past_action_id = model.action_readout.sample(logits)
+
+            # get log prob from meta controller latent actions
 
-    for one_state in state.unbind(dim = 1):
-        one_state = rearrange(one_state, 'b d -> b 1 d')
+            meta_output = cache.prev_hiddens.meta_controller
 
-        logits, cache = model(one_state, past_action_id, meta_controller = meta_controller, return_cache = True)
+            old_log_probs = meta_controller.log_prob(meta_output.action_dist, meta_output.actions)
+
+            states.append(meta_output.input_residual_stream)
+            log_probs.append(old_log_probs)
+            switch_betas.append(meta_output.switch_beta)
+            latent_actions.append(meta_output.actions)
+
+        # accumulate across time for the episode data
+
+        all_episodes.append(dict(
+            states = cat(states, dim = 1),
+            log_probs = cat(log_probs, dim = 1),
+            switch_betas = cat(switch_betas, dim = 1),
+            latent_actions = cat(latent_actions, dim = 1)
+        ))
+
+        all_rewards.append(torch.randn(1))
+
+    # calculate advantages using z-score
+
+    rewards = cat(all_rewards)
+    advantages = z_score(rewards)
+
+    assert advantages.shape == (3,)
+
+    # simulate a policy loss update over the entire group
+
+    group_states = cat([e['states'] for e in all_episodes], dim = 0)
+    group_log_probs = cat([e['log_probs'] for e in all_episodes], dim = 0)
+    group_latent_actions = cat([e['latent_actions'] for e in all_episodes], dim = 0)
+    group_switch_betas = cat([e['switch_betas'] for e in all_episodes], dim = 0)
+
+    if not use_binary_mapper_variant:
+        loss = policy_loss(
+            meta_controller,
+            group_states,
+            group_log_probs,
+            group_latent_actions,
+            advantages,
+            group_switch_betas == 1.,
+            episode_lens = episode_lens[:1].repeat(3) if exists(episode_lens) else None
+        )
 
-        assert logits.shape == (2, 1, *assert_shape)
-        past_action_id = model.action_readout.sample(logits)
+        loss.backward()
 
     # evolutionary strategies over grpo
 

{metacontroller_pytorch-0.0.31 → metacontroller_pytorch-0.0.33}/train_behavior_clone_babyai.py

@@ -26,7 +26,7 @@ from memmap_replay_buffer import ReplayBuffer
 from einops import rearrange
 
 from metacontroller.metacontroller import Transformer
-from metacontroller.metacontroller_with_resnet import TransformerWithResnetEncoder
+from metacontroller.transformer_with_resnet import TransformerWithResnet
 
 import minigrid
 import gymnasium as gym
@@ -95,7 +95,7 @@ def train(
 
     # transformer
     
-    transformer_class = TransformerWithResnetEncoder if use_resnet else Transformer
+    transformer_class = TransformerWithResnet if use_resnet else Transformer
     model = transformer_class(
         dim = dim,
         state_embed_readout = dict(num_continuous = state_dim),

metacontroller_pytorch-0.0.31/metacontroller/metacontroller_with_resnet.py

@@ -1,250 +0,0 @@
-from typing import Any, List, Type, Union, Optional
-
-import torch
-from torch import Tensor
-from torch import nn
-from einops import rearrange
-from metacontroller.metacontroller import Transformer
-
-class TransformerWithResnetEncoder(Transformer):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self.resnet_dim = kwargs["state_embed_readout"]["num_continuous"]
-        self.visual_encoder = resnet18(out_dim=self.resnet_dim)
-
-    def visual_encode(self, x: torch.Tensor) -> torch.Tensor:
-        b, t = x.shape[:2]
-        x = rearrange(x, 'b t h w c -> (b t) c h w')
-        h = self.visual_encoder(x)
-        h = rearrange(h, '(b t) d -> b t d', b=b, t=t, d = self.resnet_dim)
-        return h
-
-# resnet components taken from https://github.com/Lornatang/ResNet-PyTorch
-
-class _BasicBlock(nn.Module):
-    expansion: int = 1
-
-    def __init__(
-            self,
-            in_channels: int,
-            out_channels: int,
-            stride: int,
-            downsample: Optional[nn.Module] = None,
-            groups: int = 1,
-            base_channels: int = 64,
-    ) -> None:
-        super(_BasicBlock, self).__init__()
-        self.stride = stride
-        self.downsample = downsample
-        self.groups = groups
-        self.base_channels = base_channels
-
-        self.conv1 = nn.Conv2d(in_channels, out_channels, (3, 3), (stride, stride), (1, 1), bias=False)
-        self.bn1 = nn.BatchNorm2d(out_channels)
-        self.relu = nn.ReLU(True)
-        self.conv2 = nn.Conv2d(out_channels, out_channels, (3, 3), (1, 1), (1, 1), bias=False)
-        self.bn2 = nn.BatchNorm2d(out_channels)
-
-    def forward(self, x: Tensor) -> Tensor:
-        identity = x
-
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-
-        if self.downsample is not None:
-            identity = self.downsample(x)
-
-        out = torch.add(out, identity)
-        out = self.relu(out)
-
-        return out
-
-
-class _Bottleneck(nn.Module):
-    expansion: int = 4
-
-    def __init__(
-            self,
-            in_channels: int,
-            out_channels: int,
-            stride: int,
-            downsample: Optional[nn.Module] = None,
-            groups: int = 1,
-            base_channels: int = 64,
-    ) -> None:
-        super(_Bottleneck, self).__init__()
-        self.stride = stride
-        self.downsample = downsample
-        self.groups = groups
-        self.base_channels = base_channels
-
-        channels = int(out_channels * (base_channels / 64.0)) * groups
-
-        self.conv1 = nn.Conv2d(in_channels, channels, (1, 1), (1, 1), (0, 0), bias=False)
-        self.bn1 = nn.BatchNorm2d(channels)
-        self.conv2 = nn.Conv2d(channels, channels, (3, 3), (stride, stride), (1, 1), groups=groups, bias=False)
-        self.bn2 = nn.BatchNorm2d(channels)
-        self.conv3 = nn.Conv2d(channels, int(out_channels * self.expansion), (1, 1), (1, 1), (0, 0), bias=False)
-        self.bn3 = nn.BatchNorm2d(int(out_channels * self.expansion))
-        self.relu = nn.ReLU(True)
-
-    def forward(self, x: Tensor) -> Tensor:
-        identity = x
-
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-        out = self.relu(out)
-
-        out = self.conv3(out)
-        out = self.bn3(out)
-
-        if self.downsample is not None:
-            identity = self.downsample(x)
-
-        out = torch.add(out, identity)
-        out = self.relu(out)
-
-        return out
-
-
-class ResNet(nn.Module):
-
-    def __init__(
-            self,
-            arch_cfg: List[int],
-            block: Type[Union[_BasicBlock, _Bottleneck]],
-            groups: int = 1,
-            channels_per_group: int = 64,
-            out_dim: int = 1000,
-    ) -> None:
-        super(ResNet, self).__init__()
-        self.in_channels = 64
-        self.dilation = 1
-        self.groups = groups
-        self.base_channels = channels_per_group
-
-        self.conv1 = nn.Conv2d(3, self.in_channels, (7, 7), (2, 2), (3, 3), bias=False)
-        self.bn1 = nn.BatchNorm2d(self.in_channels)
-        self.relu = nn.ReLU(True)
-        self.maxpool = nn.MaxPool2d((3, 3), (2, 2), (1, 1))
-
-        self.layer1 = self._make_layer(arch_cfg[0], block, 64, 1)
-        self.layer2 = self._make_layer(arch_cfg[1], block, 128, 2)
-        self.layer3 = self._make_layer(arch_cfg[2], block, 256, 2)
-        self.layer4 = self._make_layer(arch_cfg[3], block, 512, 2)
-
-        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
-
-        self.fc = nn.Linear(512 * block.expansion, out_dim)
-
-        # Initialize neural network weights
-        self._initialize_weights()
-
-    def _make_layer(
-            self,
-            repeat_times: int,
-            block: Type[Union[_BasicBlock, _Bottleneck]],
-            channels: int,
-            stride: int = 1,
-    ) -> nn.Sequential:
-        downsample = None
-
-        if stride != 1 or self.in_channels != channels * block.expansion:
-            downsample = nn.Sequential(
-                nn.Conv2d(self.in_channels, channels * block.expansion, (1, 1), (stride, stride), (0, 0), bias=False),
-                nn.BatchNorm2d(channels * block.expansion),
-            )
-
-        layers = [
-            block(
-                self.in_channels,
-                channels,
-                stride,
-                downsample,
-                self.groups,
-                self.base_channels
-            )
-        ]
-        self.in_channels = channels * block.expansion
-        for _ in range(1, repeat_times):
-            layers.append(
-                block(
-                    self.in_channels,
-                    channels,
-                    1,
-                    None,
-                    self.groups,
-                    self.base_channels,
-                )
-            )
-
-        return nn.Sequential(*layers)
-
-    def forward(self, x: Tensor) -> Tensor:
-        out = self._forward_impl(x)
-
-        return out
-
-    # Support torch.script function
-    def _forward_impl(self, x: Tensor) -> Tensor:
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-        out = self.maxpool(out)
-
-        out = self.layer1(out)
-        out = self.layer2(out)
-        out = self.layer3(out)
-        out = self.layer4(out)
-
-        out = self.avgpool(out)
-        out = torch.flatten(out, 1)
-        out = self.fc(out)
-
-        return out
-
-    def _initialize_weights(self) -> None:
-        for module in self.modules():
-            if isinstance(module, nn.Conv2d):
-                nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
-            elif isinstance(module, (nn.BatchNorm2d, nn.GroupNorm)):
-                nn.init.constant_(module.weight, 1)
-                nn.init.constant_(module.bias, 0)
-
-
-def resnet18(**kwargs: Any) -> ResNet:
-    model = ResNet([2, 2, 2, 2], _BasicBlock, **kwargs)
-
-    return model
-
-
-def resnet34(**kwargs: Any) -> ResNet:
-    model = ResNet([3, 4, 6, 3], _BasicBlock, **kwargs)
-
-    return model
-
-
-def resnet50(**kwargs: Any) -> ResNet:
-    model = ResNet([3, 4, 6, 3], _Bottleneck, **kwargs)
-
-    return model
-
-
-def resnet101(**kwargs: Any) -> ResNet:
-    model = ResNet([3, 4, 23, 3], _Bottleneck, **kwargs)
-
-    return model
-
-
-def resnet152(**kwargs: Any) -> ResNet:
-    model = ResNet([3, 8, 36, 3], _Bottleneck, **kwargs)
-
-    return model