remove autocast

src/model/big_modules.py

@@ -371,35 +371,34 @@ class MaskDecoder(nn.Module):
             p4 = self.up_8_4(p8, f4)
             p2 = self.up_4_2(p4, f2)
             p1 = self.up_2_1(p2, f1)
-            with torch.cuda.amp.autocast(enabled=False):
-                if seg_pass:
-                    if last_mask is not None:
-                        res = self.pred_seg(
-                            F.relu(p1.flatten(start_dim=0, end_dim=1).float())
-                        )
-                        if sigmoid_residual:
-                            res = (
-                                torch.sigmoid(res) - 0.5
-                            ) * 2  # regularization: (-1, 1) change on last mask
-                        logits = last_mask + res
-                    else:
-                        logits = self.pred_seg(
-                            F.relu(p1.flatten(start_dim=0, end_dim=1).float())
-                        )
+            if seg_pass:
+                if last_mask is not None:
+                    res = self.pred_seg(
+                        F.relu(p1.flatten(start_dim=0, end_dim=1).float())
+                    )
+                    if sigmoid_residual:
+                        res = (
+                            torch.sigmoid(res) - 0.5
+                        ) * 2  # regularization: (-1, 1) change on last mask
+                    logits = last_mask + res
                 else:
-                    if last_mask is not None:
-                        res = self.pred_mat(
-                            F.relu(p1.flatten(start_dim=0, end_dim=1).float())
-                        )
-                        if sigmoid_residual:
-                            res = (
-                                torch.sigmoid(res) - 0.5
-                            ) * 2  # regularization: (-1, 1) change on last mask
-                        logits = last_mask + res
-                    else:
-                        logits = self.pred_mat(
-                            F.relu(p1.flatten(start_dim=0, end_dim=1).float())
-                        )
+                    logits = self.pred_seg(
+                        F.relu(p1.flatten(start_dim=0, end_dim=1).float())
+                    )
+            else:
+                if last_mask is not None:
+                    res = self.pred_mat(
+                        F.relu(p1.flatten(start_dim=0, end_dim=1).float())
+                    )
+                    if sigmoid_residual:
+                        res = (
+                            torch.sigmoid(res) - 0.5
+                        ) * 2  # regularization: (-1, 1) change on last mask
+                    logits = last_mask + res
+                else:
+                    logits = self.pred_mat(
+                        F.relu(p1.flatten(start_dim=0, end_dim=1).float())
+                    )
             ## SensoryUpdater_fullscale
             if update_sensory:
                 p1 = torch.cat(

src/model/matanyone.py

@@ -161,33 +161,32 @@ class MatAnyone(nn.Module):
         uncert_mask = uncert_output["mask"] if uncert_output is not None else None
 
         # read using visual attention
-        with torch.cuda.amp.autocast(enabled=False):
-            affinity = get_affinity(
-                memory_key.float(),
-                memory_shrinkage.float(),
-                query_key.float(),
-                query_selection.float(),
-                uncert_mask=uncert_mask,
-            )
+        affinity = get_affinity(
+            memory_key.float(),
+            memory_shrinkage.float(),
+            query_key.float(),
+            query_selection.float(),
+            uncert_mask=uncert_mask,
+        )
 
-            msk_value = msk_value.flatten(start_dim=1, end_dim=2).float()
+        msk_value = msk_value.flatten(start_dim=1, end_dim=2).float()
 
-            # B * (num_objects*CV) * H * W
-            pixel_readout = readout(affinity, msk_value, uncert_mask)
-            pixel_readout = pixel_readout.view(
-                batch_size, num_objects, self.value_dim, *pixel_readout.shape[-2:]
-            )
+        # B * (num_objects*CV) * H * W
+        pixel_readout = readout(affinity, msk_value, uncert_mask)
+        pixel_readout = pixel_readout.view(
+            batch_size, num_objects, self.value_dim, *pixel_readout.shape[-2:]
+        )
 
-            uncert_output = self.pred_uncertainty(
-                last_pix_feat,
-                pix_feat,
-                last_pred_mask,
-                pixel_readout[:, 0] - msk_value[:, :, -1],
-            )
-            uncert_prob = uncert_output["prob"].unsqueeze(1)  # b n 1 h w
-            pixel_readout = pixel_readout * uncert_prob + msk_value[:, :, -1].unsqueeze(
-                1
-            ) * (1 - uncert_prob)
+        uncert_output = self.pred_uncertainty(
+            last_pix_feat,
+            pix_feat,
+            last_pred_mask,
+            pixel_readout[:, 0] - msk_value[:, :, -1],
+        )
+        uncert_prob = uncert_output["prob"].unsqueeze(1)  # b n 1 h w
+        pixel_readout = pixel_readout * uncert_prob + msk_value[:, :, -1].unsqueeze(
+            1
+        ) * (1 - uncert_prob)
 
         pixel_readout = self.pixel_fusion(pix_feat, pixel_readout, sensory, last_mask)
 

src/model/modules.py

@@ -94,11 +94,10 @@ class SensoryUpdater_fullscale(nn.Module):
             + self.g1_conv(downsample_groups(g[4], ratio=1 / 16))
         )
 
-        with torch.cuda.amp.autocast(enabled=False):
-            g = g.float()
-            h = h.float()
-            values = self.transform(torch.cat([g, h], dim=2))
-            new_h = _recurrent_update(h, values)
+        g = g.float()
+        h = h.float()
+        values = self.transform(torch.cat([g, h], dim=2))
+        new_h = _recurrent_update(h, values)
 
         return new_h
 
@@ -124,11 +123,10 @@ class SensoryUpdater(nn.Module):
             + self.g4_conv(downsample_groups(g[2], ratio=1 / 4))
         )
 
-        with torch.cuda.amp.autocast(enabled=False):
-            g = g.float()
-            h = h.float()
-            values = self.transform(torch.cat([g, h], dim=2))
-            new_h = _recurrent_update(h, values)
+        g = g.float()
+        h = h.float()
+        values = self.transform(torch.cat([g, h], dim=2))
+        new_h = _recurrent_update(h, values)
 
         return new_h
 
@@ -143,11 +141,10 @@ class SensoryDeepUpdater(nn.Module):
         nn.init.xavier_normal_(self.transform.weight)
 
     def forward(self, g: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
-        with torch.cuda.amp.autocast(enabled=False):
-            g = g.float()
-            h = h.float()
-            values = self.transform(torch.cat([g, h], dim=2))
-            new_h = _recurrent_update(h, values)
+        g = g.float()
+        h = h.float()
+        values = self.transform(torch.cat([g, h], dim=2))
+        new_h = _recurrent_update(h, values)
 
         return new_h
 

src/model/object_summarizer.py

@@ -80,11 +80,10 @@ class ObjectSummarizer(nn.Module):
             pe = self.pos_enc(value)
             value = value + pe
 
-        with torch.cuda.amp.autocast(enabled=False):
-            value = value.float()
-            feature = self.feature_pred(value)
-            logits = self.weights_pred(value)
-            sums, area = _weighted_pooling(repeated_masks, feature, logits)
+        value = value.float()
+        feature = self.feature_pred(value)
+        logits = self.weights_pred(value)
+        sums, area = _weighted_pooling(repeated_masks, feature, logits)
 
         summaries = torch.cat([sums, area], dim=-1)
 

src/model/tensor_utils.py

@@ -46,14 +46,13 @@ def unpad(img: torch.Tensor, pad: Iterable[int]) -> torch.Tensor:
 
 # @torch.jit.script
 def aggregate(prob: torch.Tensor, dim: int) -> torch.Tensor:
-    with torch.cuda.amp.autocast(enabled=False):
-        prob = prob.float()
-        new_prob = torch.cat(
-            [torch.prod(1 - prob, dim=dim, keepdim=True), prob], dim
-        ).clamp(1e-7, 1 - 1e-7)
-        logits = torch.log((new_prob / (1 - new_prob)))  # (0, 1) --> (-inf, inf)
+    prob = prob.float()
+    new_prob = torch.cat(
+        [torch.prod(1 - prob, dim=dim, keepdim=True), prob], dim
+    ).clamp(1e-7, 1 - 1e-7)
+    logits = torch.log((new_prob / (1 - new_prob)))  # (0, 1) --> (-inf, inf)
 
-        return logits
+    return logits
 
 
 # @torch.jit.script