diff --git a/src/model/big_modules.py b/src/model/big_modules.py index a3aa5db..2c0ba94 100644 --- a/src/model/big_modules.py +++ b/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( diff --git a/src/model/matanyone.py b/src/model/matanyone.py index 7efc45b..1867d86 100644 --- a/src/model/matanyone.py +++ b/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) diff --git a/src/model/modules.py b/src/model/modules.py index 3bcabf5..881b4eb 100644 --- a/src/model/modules.py +++ b/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 diff --git a/src/model/object_summarizer.py b/src/model/object_summarizer.py index 1d6c200..fb64be2 100644 --- a/src/model/object_summarizer.py +++ b/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) diff --git a/src/model/tensor_utils.py b/src/model/tensor_utils.py index 2b69979..028bfa0 100644 --- a/src/model/tensor_utils.py +++ b/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