update: roma

third_party/Roma/roma/models/encoders.py

@@ -38,10 +38,13 @@ class ResNet50(nn.Module):
         self.freeze_bn = freeze_bn
         self.early_exit = early_exit
         self.amp = amp
-        if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
-            self.amp_dtype = torch.bfloat16
+        if torch.cuda.is_available():
+            if torch.cuda.is_bf16_supported():
+                self.amp_dtype = torch.bfloat16
+            else:
+                self.amp_dtype = torch.float16
         else:
-            self.amp_dtype = torch.float16
+            self.amp_dtype = torch.float32
 
     def forward(self, x, **kwargs):
         with torch.autocast("cuda", enabled=self.amp, dtype=self.amp_dtype):
@@ -78,10 +81,13 @@ class VGG19(nn.Module):
         super().__init__()
         self.layers = nn.ModuleList(tvm.vgg19_bn(pretrained=pretrained).features[:40])
         self.amp = amp
-        if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
-            self.amp_dtype = torch.bfloat16
+        if torch.cuda.is_available():
+            if torch.cuda.is_bf16_supported():
+                self.amp_dtype = torch.bfloat16
+            else:
+                self.amp_dtype = torch.float16
         else:
-            self.amp_dtype = torch.float16
+            self.amp_dtype = torch.float32
 
     def forward(self, x, **kwargs):
         with torch.autocast("cuda", enabled=self.amp, dtype=self.amp_dtype):
@@ -121,10 +127,13 @@ class CNNandDinov2(nn.Module):
         else:
             self.cnn = VGG19(**cnn_kwargs)
         self.amp = amp
-        if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
-            self.amp_dtype = torch.bfloat16
+        if torch.cuda.is_available():
+            if torch.cuda.is_bf16_supported():
+                self.amp_dtype = torch.bfloat16
+            else:
+                self.amp_dtype = torch.float16
         else:
-            self.amp_dtype = torch.float16
+            self.amp_dtype = torch.float32
         if self.amp:
             dinov2_vitl14 = dinov2_vitl14.to(self.amp_dtype)
         self.dinov2_vitl14 = [dinov2_vitl14]  # ugly hack to not show parameters to DDP

third_party/Roma/roma/models/matcher.py

@@ -76,10 +76,13 @@ class ConvRefiner(nn.Module):
         self.disable_local_corr_grad = disable_local_corr_grad
         self.is_classifier = is_classifier
         self.sample_mode = sample_mode
-        if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
-            self.amp_dtype = torch.bfloat16
+        if torch.cuda.is_available():
+            if torch.cuda.is_bf16_supported():
+                self.amp_dtype = torch.bfloat16
+            else:
+                self.amp_dtype = torch.float16
         else:
-            self.amp_dtype = torch.float16
+            self.amp_dtype = torch.float32
 
     def create_block(
         self,
@@ -337,10 +340,13 @@ class Decoder(nn.Module):
         self.displacement_dropout_p = displacement_dropout_p
         self.gm_warp_dropout_p = gm_warp_dropout_p
         self.flow_upsample_mode = flow_upsample_mode
-        if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
-            self.amp_dtype = torch.bfloat16
+        if torch.cuda.is_available():
+            if torch.cuda.is_bf16_supported():
+                self.amp_dtype = torch.bfloat16
+            else:
+                self.amp_dtype = torch.float16
         else:
-            self.amp_dtype = torch.float16
+            self.amp_dtype = torch.float32
 
     def get_placeholder_flow(self, b, h, w, device):
         coarse_coords = torch.meshgrid(

third_party/Roma/roma/models/transformer/__init__.py

@@ -30,10 +30,14 @@ class TransformerDecoder(nn.Module):
         self._scales = [16]
         self.is_classifier = is_classifier
         self.amp = amp
-        if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
-            self.amp_dtype = torch.bfloat16
+        if torch.cuda.is_available():
+            if torch.cuda.is_bf16_supported():
+                self.amp_dtype = torch.bfloat16
+            else:
+                self.amp_dtype = torch.float16
         else:
-            self.amp_dtype = torch.float16
+            self.amp_dtype = torch.float32
+
         self.pos_enc = pos_enc
         self.learned_embeddings = learned_embeddings
         if self.learned_embeddings: