add profiling

unik3d/layers/upsample.py

@@ -89,7 +89,7 @@ class ResUpsampleBil(nn.Module):
             nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
         )
 
-    @profile_method(verbose=VERBOSE)
+    @profile_method(verbose=True)
     def forward(self, x: torch.Tensor):
         for conv in self.convs:
             x = conv(x)
@@ -122,7 +122,7 @@ class ResUpsample(nn.Module):
             hidden_dim, hidden_dim // 2, kernel_size=2, stride=2, padding=0
         )
 
-    @profile_method(verbose=VERBOSE)
+    @profile_method(verbose=True)
     def forward(self, x: torch.Tensor):
         for conv in self.convs:
             x = conv(x)

unik3d/models/decoder.py

@@ -18,7 +18,7 @@ from unik3d.utils.geometric import flat_interpolate
 from unik3d.utils.misc import get_params
 from unik3d.utils.positional_embedding import generate_fourier_features
 from unik3d.utils.sht import rsh_cart_3
-
+from unik3d.utils.misc import profile_method
 
 def orthonormal_init(num_tokens, dims):
     pe = torch.randn(num_tokens, dims)
@@ -120,6 +120,7 @@ class AngularModule(nn.Module):
         intrinsics = correction_tensor.unsqueeze(0) * intrinsics
         return intrinsics
 
+    @profile_method(verbose=True)
     def forward(self, cls_tokens) -> torch.Tensor:
         latents_pos = self.latents_pos.expand(cls_tokens.shape[0], -1, -1)
 
@@ -271,6 +272,7 @@ class RadialModule(nn.Module):
     def set_shapes(self, shapes: tuple[int, int]):
         self.shapes = shapes
 
+    @profile_method(verbose=True)
     def embed_rays(self, rays):
         rays_embedding = flat_interpolate(
             rays, old=self.original_shapes, new=self.shapes, antialias=True
@@ -292,6 +294,7 @@ class RadialModule(nn.Module):
         )
         return rays_embedding
 
+    @profile_method(verbose=True)
     def condition(self, feat, rays_embeddings):
         conditioned_features = [
             prompter(rearrange(feature, "b h w c -> b (h w) c"), rays_embeddings)
@@ -321,6 +324,7 @@ class RadialModule(nn.Module):
 
         return out_features, init_latents
 
+    @profile_method(verbose=True)
     def depth_proj(self, out_features):
         depths = []
         h_out, w_out = out_features[-1].shape[-2:]
@@ -340,6 +344,7 @@ class RadialModule(nn.Module):
         logdepth = self.to_depth_hr(logdepth)
         return logdepth
 
+    @profile_method(verbose=True)
     def confidence_proj(self, out_features):
         highres_features = out_features[-1].permute(0, 2, 3, 1)
         confidence = self.confidence_mlp(highres_features).permute(0, 3, 1, 2)

unik3d/models/metadinov2/block.py

@@ -21,6 +21,7 @@ from .mlp import Mlp
 
 logger = logging.getLogger("dinov2")
 
+from unik3d.utils.misc import profile_method
 
 try:
     from xformers.ops import fmha, index_select_cat, scaled_index_add
@@ -80,6 +81,7 @@ class Block(nn.Module):
 
         self.sample_drop_ratio = drop_path
 
+    @profile_method(verbose=True)
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         def attn_residual_func(x: torch.Tensor) -> torch.Tensor:
             return self.ls1(self.attn(self.norm1(x)))