(args)
| 616 | |
| 617 | |
| 618 | def build(args): |
| 619 | dataset_to_num_classes = { |
| 620 | 'coco': 91, |
| 621 | 'coco_panoptic': 250, |
| 622 | 'e2e_mot': 1, |
| 623 | 'e2e_joint': 1, |
| 624 | 'e2e_static_mot': 1 |
| 625 | } |
| 626 | assert args.dataset_file in dataset_to_num_classes |
| 627 | num_classes = dataset_to_num_classes[args.dataset_file] |
| 628 | device = torch.device(args.device) |
| 629 | |
| 630 | backbone = build_backbone(args) |
| 631 | |
| 632 | transformer = build_deforamble_transformer(args) |
| 633 | d_model = transformer.d_model |
| 634 | hidden_dim = args.dim_feedforward |
| 635 | query_interaction_layer = build_query_interaction_layer(args, args.query_interaction_layer, d_model, hidden_dim, d_model*2) |
| 636 | |
| 637 | img_matcher = build_matcher(args) |
| 638 | num_frames_per_batch = max(args.sampler_lengths) |
| 639 | weight_dict = {} |
| 640 | for i in range(num_frames_per_batch): |
| 641 | weight_dict.update({"frame_{}_loss_ce".format(i): args.cls_loss_coef, |
| 642 | 'frame_{}_loss_bbox'.format(i): args.bbox_loss_coef, |
| 643 | 'frame_{}_loss_giou'.format(i): args.giou_loss_coef, |
| 644 | }) |
| 645 | |
| 646 | # TODO this is a hack |
| 647 | if args.aux_loss: |
| 648 | for i in range(num_frames_per_batch): |
| 649 | for j in range(args.dec_layers - 1): |
| 650 | weight_dict.update({"frame_{}_aux{}_loss_ce".format(i, j): args.cls_loss_coef, |
| 651 | 'frame_{}_aux{}_loss_bbox'.format(i, j): args.bbox_loss_coef, |
| 652 | 'frame_{}_aux{}_loss_giou'.format(i, j): args.giou_loss_coef, |
| 653 | }) |
| 654 | if args.memory_bank_type is not None and len(args.memory_bank_type) > 0: |
| 655 | memory_bank = build_memory_bank(args, d_model, hidden_dim, d_model * 2) |
| 656 | for i in range(num_frames_per_batch): |
| 657 | weight_dict.update({"frame_{}_track_loss_ce".format(i): args.cls_loss_coef}) |
| 658 | else: |
| 659 | memory_bank = None |
| 660 | losses = ['labels', 'boxes'] |
| 661 | criterion = ClipMatcher(num_classes, matcher=img_matcher, weight_dict=weight_dict, losses=losses) |
| 662 | criterion.to(device) |
| 663 | postprocessors = {} |
| 664 | model = MOTR( |
| 665 | backbone, |
| 666 | transformer, |
| 667 | track_embed=query_interaction_layer, |
| 668 | num_feature_levels=args.num_feature_levels, |
| 669 | num_classes=num_classes, |
| 670 | num_queries=args.num_queries, |
| 671 | aux_loss=args.aux_loss, |
| 672 | criterion=criterion, |
| 673 | with_box_refine=args.with_box_refine, |
| 674 | two_stage=args.two_stage, |
| 675 | memory_bank=memory_bank, |
nothing calls this directly
no test coverage detected