ef35a38ad29c@2024-04-07_09-03-46: fix ape-e with fsdp

ape/engine/defaults.py

@@ -108,6 +108,8 @@ def create_fsdp_model(model, *, fp16_compression=False, **kwargs):
  # _module_classes_to_ignore=(MultiScaleDeformableAttention,),
  )
 
+ model = model.to(param_dtype)
+
  fsdp = FSDP(
  model,
  auto_wrap_policy=auto_wrap_policy,

ape/evaluation/oideval.py

@@ -53,7 +53,7 @@ def compute_average_precision(precision, recall):
 
  if not isinstance(precision, np.ndarray) or not isinstance(recall, np.ndarray):
  raise ValueError("precision and recall must be numpy array")
- if precision.dtype != np.float or recall.dtype != np.float:
+ if precision.dtype != float or recall.dtype != float:
  raise ValueError("input must be float numpy array.")
  if len(precision) != len(recall):
  raise ValueError("precision and recall must be of the same size.")
@@ -448,8 +448,8 @@ def accumulate(self):
  tps = np.logical_and(dt_m, np.logical_not(dt_ig))
  fps = np.logical_and(np.logical_not(dt_m), np.logical_not(dt_ig))
 
- tp_sum = np.cumsum(tps, axis=1).astype(dtype=np.float)
- fp_sum = np.cumsum(fps, axis=1).astype(dtype=np.float)
+ tp_sum = np.cumsum(tps, axis=1).astype(dtype=float)
+ fp_sum = np.cumsum(fps, axis=1).astype(dtype=float)
 
  dt_pointers[cat_idx][area_idx] = {
  "tps": tps,
@@ -479,7 +479,7 @@ def accumulate(self):
  pr[i - 1] = pr[i]
 
  mAP = compute_average_precision(
- np.array(pr, np.float).reshape(-1), np.array(rc, np.float).reshape(-1)
+ np.array(pr, float).reshape(-1), np.array(rc, float).reshape(-1)
  )
  precision[iou_thr_idx, :, cat_idx, area_idx] = mAP
 

ape/layers/fuse_helper.py

@@ -219,8 +219,10 @@ def __init__(
  self.gamma_l = nn.Parameter(init_values * torch.ones((l_dim)), requires_grad=True)
 
  def forward(self, v, l, attention_mask_v=None, attention_mask_l=None):
- v = self.layer_norm_v(v.float())
- l = self.layer_norm_l(l.float())
+ # v = self.layer_norm_v(v.float())
+ # l = self.layer_norm_l(l.float())
+ v = self.layer_norm_v(v)
+ l = self.layer_norm_l(l)
  delta_v, delta_l = self.attn(
  v, l, attention_mask_v=attention_mask_v, attention_mask_l=attention_mask_l
  )

ape/layers/vision_language_fusion.py

@@ -45,7 +45,7 @@ def __init__(
 
  def forward(self, v, l, attention_mask_v=None, attention_mask_l=None):
  if self.use_checkpoint and self.training:
- return checkpoint.checkpoint(self.b_attn, v, l, attention_mask_v, attention_mask_l)
+ return checkpoint.checkpoint(self.b_attn, v, l, attention_mask_v, attention_mask_l, use_reentrant=False)
  else:
  return self.b_attn(v, l, attention_mask_v, attention_mask_l)
 

ape/modeling/ape_deta/deformable_detr.py

@@ -76,6 +76,7 @@ def __init__(
  text_feature_batch_repeat: bool = True,
  text_feature_bank: bool = False,
  text_feature_bank_reset: bool = False,
+ text_feature_bank_random_size: bool = False,
  text_feature_reduce_type: str = "last",
  text_feature_reduce_before_fusion: bool = True,
  expression_cumulative_gt_class: bool = True,
@@ -276,6 +277,7 @@ def __init__(
 
  self.text_feature_bank = text_feature_bank
  self.text_feature_bank_reset = text_feature_bank_reset
+ self.text_feature_bank_random_size = text_feature_bank_random_size
  if self.text_feature_bank:
  features_phrase_bank = torch.zeros(
  (

ape/modeling/ape_deta/deformable_detr_segm.py

@@ -324,6 +324,11 @@ def forward(self, batched_inputs, do_postprocess=True):
  : max(len(text_list), self.criterion[dataset_id].num_classes)
  ]
 
+ if self.text_feature_bank and self.text_feature_bank_random_size:
+ features_l = features_l[
+ : random.randint(len(text_list), len(features_l))
+ ]
+
  if self.text_feature_batch_repeat:
  features_l = features_l.unsqueeze(0).repeat(len(batched_inputs), 1, 1)
  else:

ape/modeling/ape_deta/deformable_detr_segm_vl.py

@@ -1,4 +1,5 @@
 import copy
+import random
 import math
 import os
 import time
@@ -324,6 +325,12 @@ def forward(self, batched_inputs, do_postprocess=True):
  : max(len(text_list), self.criterion[dataset_id].num_classes)
  ]
 
+ if self.text_feature_bank and self.text_feature_bank_random_size:
+ text_feature_bank_size = random.randint(len(text_list), len(features_l))
+ features_l = features_l[
+ : random.randint(len(text_list), len(features_l))
+ ]
+
  if self.text_feature_batch_repeat:
  features_l = features_l.unsqueeze(0).repeat(len(batched_inputs), 1, 1)
  else:

ape/modeling/ape_deta/deformable_transformer.py

@@ -2,6 +2,7 @@
 
 import torch
 import torch.nn as nn
+import torch.utils.checkpoint as checkpoint
 
 from ape.layers import MultiScaleDeformableAttention
 from detrex.layers import (
@@ -59,12 +60,7 @@ def __init__(
  else:
  self.post_norm_layer = None
 
- if use_act_checkpoint:
- from fairscale.nn.checkpoint import checkpoint_wrapper
-
- for i, layer in enumerate(self.layers):
- layer = checkpoint_wrapper(layer)
- self.layers[i] = layer
+ self.use_checkpoint = use_act_checkpoint
 
  def forward(
  self,
@@ -80,16 +76,30 @@ def forward(
  ):
 
  for layer in self.layers:
- query = layer(
- query,
- key,
- value,
- query_pos=query_pos,
- attn_masks=attn_masks,
- query_key_padding_mask=query_key_padding_mask,
- key_padding_mask=key_padding_mask,
- **kwargs,
- )
+ if self.use_checkpoint and self.training:
+ query = checkpoint.checkpoint(
+ layer,
+ query,
+ key,
+ value,
+ query_pos=query_pos,
+ attn_masks=attn_masks,
+ query_key_padding_mask=query_key_padding_mask,
+ key_padding_mask=key_padding_mask,
+ use_reentrant=False,
+ **kwargs,
+ )
+ else:
+ query = layer(
+ query,
+ key,
+ value,
+ query_pos=query_pos,
+ attn_masks=attn_masks,
+ query_key_padding_mask=query_key_padding_mask,
+ key_padding_mask=key_padding_mask,
+ **kwargs,
+ )
 
  if self.post_norm_layer is not None:
  query = self.post_norm_layer(query)
@@ -144,12 +154,7 @@ def __init__(
  self.bbox_embed = None
  self.class_embed = None
 
- if use_act_checkpoint:
- from fairscale.nn.checkpoint import checkpoint_wrapper
-
- for i, layer in enumerate(self.layers):
- layer = checkpoint_wrapper(layer)
- self.layers[i] = layer
+ self.use_checkpoint = use_act_checkpoint
 
  def forward(
  self,
@@ -179,18 +184,34 @@ def forward(
  assert reference_points.shape[-1] == 2
  reference_points_input = reference_points[:, :, None] * valid_ratios[:, None]
 
- output = layer(
- output,
- key,
- value,
- query_pos=query_pos,
- key_pos=key_pos,
- attn_masks=attn_masks,
- query_key_padding_mask=query_key_padding_mask,
- key_padding_mask=key_padding_mask,
- reference_points=reference_points_input,
- **kwargs,
- )
+ if self.use_checkpoint and self.training:
+ output = checkpoint.checkpoint(
+ layer,
+ output,
+ key,
+ value,
+ query_pos=query_pos,
+ key_pos=key_pos,
+ attn_masks=attn_masks,
+ query_key_padding_mask=query_key_padding_mask,
+ key_padding_mask=key_padding_mask,
+ reference_points=reference_points_input,
+ use_reentrant=False,
+ **kwargs,
+ )
+ else:
+ output = layer(
+ output,
+ key,
+ value,
+ query_pos=query_pos,
+ key_pos=key_pos,
+ attn_masks=attn_masks,
+ query_key_padding_mask=query_key_padding_mask,
+ key_padding_mask=key_padding_mask,
+ reference_points=reference_points_input,
+ **kwargs,
+ )
 
  if self.bbox_embed is not None:
  tmp = self.bbox_embed[layer_idx](output)

ape/modeling/ape_deta/deformable_transformer_vl.py

@@ -3,6 +3,7 @@
 
 import torch
 import torch.nn as nn
+import torch.utils.checkpoint as checkpoint
 
 from ape.layers import MultiScaleDeformableAttention
 from detrex.layers import (
@@ -28,7 +29,7 @@ def __init__(
  post_norm: bool = False,
  num_feature_levels: int = 4,
  vl_layer=None,
- use_act_checkpoint=False,
+ use_act_checkpoint: bool = False,
  pytorch_attn=False,
  ):
  super(DeformableDetrTransformerEncoderVL, self).__init__(
@@ -63,12 +64,7 @@ def __init__(
 
  self.vl_layers = nn.ModuleList([copy.deepcopy(vl_layer) for _ in range(num_layers)])
 
- if use_act_checkpoint:
- from fairscale.nn.checkpoint import checkpoint_wrapper
-
- for i, layer in enumerate(self.layers):
- layer = checkpoint_wrapper(layer)
- self.layers[i] = layer
+ self.use_checkpoint = use_act_checkpoint
 
  def forward(
  self,
@@ -93,16 +89,30 @@ def forward(
  attention_mask_v=query_key_padding_mask,
  attention_mask_l=attention_mask_l,
  )
- query = layer(
- query,
- key,
- value,
- query_pos=query_pos,
- attn_masks=attn_masks,
- query_key_padding_mask=query_key_padding_mask,
- key_padding_mask=key_padding_mask,
- **kwargs,
- )
+ if self.use_checkpoint and self.training:
+ query = checkpoint.checkpoint(
+ layer,
+ query,
+ key,
+ value,
+ query_pos=query_pos,
+ attn_masks=attn_masks,
+ query_key_padding_mask=query_key_padding_mask,
+ key_padding_mask=key_padding_mask,
+ use_reentrant=False,
+ **kwargs,
+ )
+ else:
+ query = layer(
+ query,
+ key,
+ value,
+ query_pos=query_pos,
+ attn_masks=attn_masks,
+ query_key_padding_mask=query_key_padding_mask,
+ key_padding_mask=key_padding_mask,
+ **kwargs,
+ )
 
  if self.post_norm_layer is not None:
  query = self.post_norm_layer(query)
@@ -160,12 +170,7 @@ def __init__(
  self.bbox_embed = None
  self.class_embed = None
 
- if use_act_checkpoint:
- from fairscale.nn.checkpoint import checkpoint_wrapper
-
- for i, layer in enumerate(self.layers):
- layer = checkpoint_wrapper(layer)
- self.layers[i] = layer
+ self.use_checkpoint = use_act_checkpoint
 
  self.look_forward_twice = look_forward_twice
 
@@ -197,18 +202,34 @@ def forward(
  assert reference_points.shape[-1] == 2
  reference_points_input = reference_points[:, :, None] * valid_ratios[:, None]
 
- output = layer(
- output,
- key,
- value,
- query_pos=query_pos,
- key_pos=key_pos,
- attn_masks=attn_masks,
- query_key_padding_mask=query_key_padding_mask,
- key_padding_mask=key_padding_mask,
- reference_points=reference_points_input,
- **kwargs,
- )
+ if self.use_checkpoint and self.training:
+ output = checkpoint.checkpoint(
+ layer,
+ output,
+ key,
+ value,
+ query_pos=query_pos,
+ key_pos=key_pos,
+ attn_masks=attn_masks,
+ query_key_padding_mask=query_key_padding_mask,
+ key_padding_mask=key_padding_mask,
+ reference_points=reference_points_input,
+ use_reentrant=False,
+ **kwargs,
+ )
+ else:
+ output = layer(
+ output,
+ key,
+ value,
+ query_pos=query_pos,
+ key_pos=key_pos,
+ attn_masks=attn_masks,
+ query_key_padding_mask=query_key_padding_mask,
+ key_padding_mask=key_padding_mask,
+ reference_points=reference_points_input,
+ **kwargs,
+ )
 
  if self.bbox_embed is not None:
  tmp = self.bbox_embed[layer_idx](output)

ape/modeling/backbone/vit_eva_clip.py

@@ -38,7 +38,9 @@ def forward(self, x: torch.Tensor):
  from apex.normalization import FusedLayerNorm
 except:
  FusedLayerNorm = LayerNorm
- print("apex.normalization.FusedLayerNorm not found, will use pytorch implementations")
+ # print("apex.normalization.FusedLayerNorm not found, will use pytorch implementations")
+
+has_sdp_kernel = hasattr(torch.backends.cuda, "sdp_kernel")
 
 
 logger = logging.getLogger(__name__)
@@ -256,7 +258,7 @@ def forward(self, x, rel_pos_bias=None, attn_mask=None):
  q = self.rope(q).type_as(v)
  k = self.rope(k).type_as(v)
 
- if True:
+ if has_sdp_kernel:
  with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True):
  x = F.scaled_dot_product_attention(q, k, v, dropout_p=self.xattn_drop, scale=self.scale)
  x = x.permute(0, 2, 1, 3) # B, num_heads, N, C -> B, N, num_heads, C