Open wm901115nwpu opened 2 years ago
def permute_preds(self, mlvl_preds): permute_mlvl_preds = [] for loc_pred, cls_pred, obj_pred in mlvl_preds: bs, ls, ny, nx = loc_pred.shape cs = cls_pred.shape[1] loc_pred = loc_pred.view(bs, self.num_anchors_per_level, ls // self.num_anchors_per_level, ny, nx).permute(0, 1, 3, 4, 2).contiguous() cls_pred = cls_pred.view(bs, self.num_anchors_per_level, cs // self.num_anchors_per_level, ny, nx).permute(0, 1, 3, 4, 2).contiguous() obj_pred = obj_pred.view(bs, self.num_anchors_per_level, 1, ny, nx).permute(0, 1, 3, 4, 2).contiguous() permute_mlvl_preds.append((loc_pred, cls_pred, obj_pred)) return permute_mlvl_preds
torch.fx.tracer proxy iterated
def permute_preds(self, mlvl_preds): permute_mlvl_preds = [] for loc_pred, cls_pred, obj_pred in mlvl_preds: bs, ls, ny, nx = loc_pred.shape cs = cls_pred.shape[1] loc_pred = loc_pred.view(bs, self.num_anchors_per_level, ls // self.num_anchors_per_level, ny, nx).permute(0, 1, 3, 4, 2).contiguous() cls_pred = cls_pred.view(bs, self.num_anchors_per_level, cs // self.num_anchors_per_level, ny, nx).permute(0, 1, 3, 4, 2).contiguous() obj_pred = obj_pred.view(bs, self.num_anchors_per_level, 1, ny, nx).permute(0, 1, 3, 4, 2).contiguous() permute_mlvl_preds.append((loc_pred, cls_pred, obj_pred)) return permute_mlvl_preds
torch.fx.tracer proxy iterated