trying to get basic model working

models/convnext.py

@@ -1,6 +1,4 @@
-# from torchvision.models.convnext
-import os
-os.environ["CUDA_VISIBLE_DEVICES"] = '1'
+# taken from torchvision.models.convnext
 
 from functools import partial
 from typing import Any, Callable, List, Optional, Sequence
@@ -11,7 +9,7 @@ from torch.nn import functional as F
 
 from torchvision.ops.misc import Conv2dNormActivation, Permute
 from torchvision.ops.stochastic_depth import StochasticDepth
-from torchvision.models._api import register_model, Weights, WeightsEnum
+from torchvision.models._api import Weights, WeightsEnum
 
 
 
@@ -163,14 +161,112 @@ class ConvNeXtBase(nn.Module):
         return self._forward_impl(x)
 
 
+
+class ConvNeXtReverseBase(nn.Module):
+    def __init__(
+        self,
+        block_setting: List[CNBlockConfig],
+        stochastic_depth_prob: float = 0.0,
+        layer_scale: float = 1e-6,
+        block: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Optional[Callable[..., nn.Module]] = None
+    ) -> None:
+        super().__init__()
+
+        if not block_setting:
+            raise ValueError("The block_setting should not be empty")
+        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
+            raise TypeError("The block_setting should be List[CNBlockConfig]")
+
+        if block is None:
+            block = CNBlock
+
+        if norm_layer is None:
+            norm_layer = partial(LayerNorm2d, eps=1e-6)
+
+        layers: List[nn.Module] = []
+
+
+        # add bottleneck which will be executed at the end, to process after the final larger upsampling
+        stage: List[nn.Module] = []
+        for _ in range(3):
+            stage.append(block(1, layer_scale, 0.0))
+        layers.append(nn.Sequential(*stage))
+
+
+        # Stem
+        firstconv_output_channels = block_setting[0].input_channels
+        layers.append(
+            # Convtr NormActivation(
+            #     firstconv_output_channels,
+            #     3,
+            #     kernel_size=4,
+            #     stride=4,
+            #     padding=0,
+            #     norm_layer=None,
+            #     activation_layer=None,
+            #     bias=True,
+            # )
+            nn.Sequential(
+                nn.ConvTranspose2d(
+                    firstconv_output_channels,
+                    1,
+                    kernel_size=4,
+                    stride=4,
+                    padding=0,
+                    ),
+            )
+        )
+
+        total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
+        stage_block_id = 0
+        for cnf in block_setting:
+            # Bottlenecks
+            stage: List[nn.Module] = []
+            for _ in range(cnf.num_layers):
+                # adjust stochastic depth probability based on the depth of the stage block
+                sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
+                stage.append(block(cnf.input_channels, layer_scale, sd_prob))
+                stage_block_id += 1
+            layers.append(nn.Sequential(*stage))
+            if cnf.out_channels is not None:
+                # Upsampling
+                layers.append(
+                    nn.Sequential(
+                        norm_layer(cnf.out_channels),
+                        nn.ConvTranspose2d(cnf.out_channels, cnf.input_channels, kernel_size=2, stride=2),
+                    )
+                )
+
+        self.features = nn.Sequential(*reversed(layers))
+
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.Linear)):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+    def _forward_impl(self, x: Tensor) -> Tensor:
+        x = self.features(x)
+        x = torch.tanh(x)
+        return x
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
+
+
 def _convnext(
     block_setting: List[CNBlockConfig],
     stochastic_depth_prob: float,
     weights: Optional[WeightsEnum],
+    reversed: bool,
     **kwargs: Any,
 ) -> ConvNeXtBase:
 
-    model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
+    if reversed:
+        model = ConvNeXtReverseBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
+    else:
+        model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
 
     if weights is not None:
         model.load_state_dict(weights.get_state_dict())
@@ -179,7 +275,7 @@ def _convnext(
 
 
 class ConvNeXt(nn.Module):
-    def __init__(self, stochastic_depth_prob=0.0, **kwargs) -> None:
+    def __init__(self, stochastic_depth_prob=0.0, reversed=False, **kwargs) -> None:
         super().__init__()
 
 
@@ -189,7 +285,7 @@ class ConvNeXt(nn.Module):
             CNBlockConfig(384, 768, 9),
             CNBlockConfig(768, None, 3),
         ]
-        self.model = _convnext(block_setting, stochastic_depth_prob, None, **kwargs)
+        self.model = _convnext(block_setting, stochastic_depth_prob, None, reversed, **kwargs)
 
     def forward(self, x: Tensor) -> Tensor:
         x = self.model(x)