Feature/profile sfno

fme/core/models/conditional_sfno/layers.py

@@ -25,6 +25,7 @@
 from torch.utils.checkpoint import checkpoint
 
 from fme.core.models.conditional_sfno.lora import LoRAConv2d
+from fme.core.models.conditional_sfno.timer import CUDATimer, NullTimer
 
 from .activations import ComplexReLU
 from .contractions import compl_mul2d_fwd, compl_muladd2d_fwd
@@ -163,35 +164,33 @@ def __init__(
             self.W_bias_labels = None
         if self.embed_dim_noise > 0:
             # no bias as it is already handled in the non-2d layers
-            self.W_scale_2d = nn.Conv2d(
-                self.embed_dim_noise, self.n_channels, kernel_size=1, bias=False
+            self.W_scale_2d = nn.Linear(
+                self.embed_dim_noise, self.n_channels, bias=False
             )
-            self.W_bias_2d = nn.Conv2d(
-                self.embed_dim_noise, self.n_channels, kernel_size=1, bias=False
+            self.W_bias_2d = nn.Linear(
+                self.embed_dim_noise, self.n_channels, bias=False
             )
         else:
             self.W_scale_2d = None
             self.W_bias_2d = None
         if self.embed_dim_pos > 0:
             # no bias as it is already handled in the non-2d layers
-            self.W_scale_pos = nn.Conv2d(
-                self.embed_dim_pos, self.n_channels, kernel_size=1, bias=False
-            )
-            self.W_bias_pos = nn.Conv2d(
-                self.embed_dim_pos, self.n_channels, kernel_size=1, bias=False
+            self.W_scale_pos = nn.Linear(
+                self.embed_dim_pos, self.n_channels, bias=False
             )
+            self.W_bias_pos = nn.Linear(self.embed_dim_pos, self.n_channels, bias=False)
         else:
             self.W_scale_pos = None
             self.W_bias_pos = None
         if global_layer_norm:
             self.norm = nn.LayerNorm(
-                (self.n_channels, img_shape[0], img_shape[1]),
+                (img_shape[1], img_shape[0], self.n_channels),
                 eps=epsilon,
                 elementwise_affine=elementwise_affine,
             )
         else:
-            self.norm = ChannelLayerNorm(
-                self.n_channels,
+            self.norm = nn.LayerNorm(
+                (self.n_channels,),
                 eps=epsilon,
                 elementwise_affine=elementwise_affine,
             )
@@ -223,7 +222,12 @@ def reset_parameters(self):
             torch.nn.init.constant_(self.W_bias_pos.weight, 0.0)
         # no bias on 2d layers as it is already handled in the non-2d layers
 
-    def forward(self, x: torch.Tensor, context: Context) -> torch.Tensor:
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: Context,
+        timer: CUDATimer | NullTimer | None = None,
+    ) -> torch.Tensor:
         """
         Conditional Layer Normalization
 
@@ -232,62 +236,74 @@ def forward(self, x: torch.Tensor, context: Context) -> torch.Tensor:
 
         Args:
             x: The input tensor to normalize, of shape
-                (batch_size, channels, height, width).
+                (batch_size, width, height, channels).
             context: The context to condition on.
 
         Returns:
-            The normalized tensor, of shape (batch_size, channels, height, width).
+            The normalized tensor, of shape (batch_size, width, height, channels).
         """
+        if timer is None:
+            timer = NullTimer()
         if context.labels is None and (
             self.W_scale_labels is not None or self.W_bias_labels is not None
         ):
             raise ValueError("labels must be provided")
-        if self.W_scale is not None:
-            if context.embedding_scalar is None:
-                raise ValueError("embedding_scalar must be provided")
-            scale: torch.Tensor = (
-                self.W_scale(context.embedding_scalar).unsqueeze(-1).unsqueeze(-1)
-            )
-        else:
-            scale = torch.ones(
-                list(x.shape[:-2]) + [1, 1], device=x.device, dtype=x.dtype
-            )
+        with timer.context("layer_norm_compute_scaling_and_bias"):
+            if self.W_scale is not None:
+                if context.embedding_scalar is None:
+                    raise ValueError("embedding_scalar must be provided")
+                scale: torch.Tensor = (
+                    self.W_scale(context.embedding_scalar).unsqueeze(-2).unsqueeze(-2)
+                )
+            else:
+                scale = torch.ones(
+                    list(x.shape[:-3]) + [1, 1, x.shape[-1]],
+                    device=x.device,
+                    dtype=x.dtype,
+                )
 
-        if self.W_scale_2d is not None:
-            if context.noise is None:
-                raise ValueError("embedding_2d must be provided")
-            scale = scale + self.W_scale_2d(context.noise)
-        if self.W_bias is not None:
-            if context.embedding_scalar is None:
-                raise ValueError("embedding_scalar must be provided")
-            bias: torch.Tensor = (
-                self.W_bias(context.embedding_scalar).unsqueeze(-1).unsqueeze(-1)
-            )
-        else:
-            bias = torch.zeros(
-                list(x.shape[:-2]) + [1, 1], device=x.device, dtype=x.dtype
-            )
+            if self.W_scale_2d is not None:
+                if context.noise is None:
+                    raise ValueError("embedding_2d must be provided")
+                scale = scale + self.W_scale_2d(context.noise)
+            if self.W_bias is not None:
+                if context.embedding_scalar is None:
+                    raise ValueError("embedding_scalar must be provided")
+                bias: torch.Tensor = (
+                    self.W_bias(context.embedding_scalar).unsqueeze(-2).unsqueeze(-2)
+                )
+            else:
+                bias = torch.zeros(
+                    list(x.shape[:-3]) + [1, 1, x.shape[-1]],
+                    device=x.device,
+                    dtype=x.dtype,
+                )
 
-        if self.W_scale_labels is not None:
-            scale = scale + self.W_scale_labels(context.labels).unsqueeze(-1).unsqueeze(
-                -1
-            )
-        if self.W_bias_labels is not None:
-            bias = bias + self.W_bias_labels(context.labels).unsqueeze(-1).unsqueeze(-1)
-        if self.W_bias_2d is not None:
-            if context.noise is None:
-                raise ValueError("embedding_2d must be provided")
-            bias = bias + self.W_bias_2d(context.noise)
-        if self.W_scale_pos is not None:
-            if context.embedding_pos is None:
-                raise ValueError("embedding_pos must be provided")
-            scale = scale + self.W_scale_pos(context.embedding_pos)
-        if self.W_bias_pos is not None:
-            if context.embedding_pos is None:
-                raise ValueError("embedding_pos must be provided")
-            bias = bias + self.W_bias_pos(context.embedding_pos)
-        x_norm: torch.Tensor = self.norm(x)
-        return x_norm * scale + bias
+            if self.W_scale_labels is not None:
+                scale = scale + self.W_scale_labels(context.labels).unsqueeze(
+                    -2
+                ).unsqueeze(-2)
+            if self.W_bias_labels is not None:
+                bias = bias + self.W_bias_labels(context.labels).unsqueeze(
+                    -2
+                ).unsqueeze(-2)
+            if self.W_bias_2d is not None:
+                if context.noise is None:
+                    raise ValueError("embedding_2d must be provided")
+                bias = bias + self.W_bias_2d(context.noise)
+            if self.W_scale_pos is not None:
+                if context.embedding_pos is None:
+                    raise ValueError("embedding_pos must be provided")
+                scale = scale + self.W_scale_pos(context.embedding_pos)
+            if self.W_bias_pos is not None:
+                if context.embedding_pos is None:
+                    raise ValueError("embedding_pos must be provided")
+                bias = bias + self.W_bias_pos(context.embedding_pos)
+        with timer.context("layer_norm_normalize"):
+            x_norm: torch.Tensor = self.norm(x)
+        with timer.context("layer_norm_apply_scaling_and_bias"):
+            return_value = x_norm * scale + bias
+        return return_value
 
 
 @torch.jit.script