More debug printing inside transformer encoder

models/GroundingDINO/transformer.py

@@ -560,7 +560,7 @@ class TransformerEncoder(nn.Module):
         """
 
         output = src
-
+        print("inside transformer encoder")
         # preparation and reshape
         if self.num_layers > 0:
             reference_points = self.get_reference_points(
@@ -591,8 +591,10 @@ class TransformerEncoder(nn.Module):
             # if output.isnan().any() or memory_text.isnan().any():
             #     if os.environ.get('IPDB_SHILONG_DEBUG', None) == 'INFO':
             #         import ipdb; ipdb.set_trace()
+            print("layer_id: " + str(layer_id))
             if self.fusion_layers:
                 if self.use_checkpoint:
+                    print("using checkpoint")
                     output, memory_text = checkpoint.checkpoint(
                         self.fusion_layers[layer_id],
                         output,
@@ -600,24 +602,30 @@ class TransformerEncoder(nn.Module):
                         key_padding_mask,
                         text_attention_mask,
                     )
+                    print("got checkpoint output")
                 else:
+                    print("not using checkpoint")
                     output, memory_text = self.fusion_layers[layer_id](
                         v=output,
                         l=memory_text,
                         attention_mask_v=key_padding_mask,
                         attention_mask_l=text_attention_mask,
                     )
+                    print("got fusion output")
 
             if self.text_layers:
+                print("getting text layers")
                 memory_text = self.text_layers[layer_id](
                     src=memory_text.transpose(0, 1),
                     src_mask=~text_self_attention_masks,  # note we use ~ for mask here
                     src_key_padding_mask=text_attention_mask,
                     pos=(pos_text.transpose(0, 1) if pos_text is not None else None),
                 ).transpose(0, 1)
+                print("got text output")
 
             # main process
             if self.use_transformer_ckpt:
+                print("use transformer ckpt")
                 output = checkpoint.checkpoint(
                     layer,
                     output,
@@ -627,7 +635,9 @@ class TransformerEncoder(nn.Module):
                     level_start_index,
                     key_padding_mask,
                 )
+                print("got output")
             else:
+                print("not use transformer ckpt")
                 output = layer(
                     src=output,
                     pos=pos,
@@ -636,6 +646,7 @@ class TransformerEncoder(nn.Module):
                     level_start_index=level_start_index,
                     key_padding_mask=key_padding_mask,
                 )
+                print("got output")
 
         return output, memory_text