Skip to content

Instantly share code, notes, and snippets.

@m-klasen
Last active November 22, 2023 19:48
Show Gist options
  • Save m-klasen/651297e28199b4bb7907fc413c49f58f to your computer and use it in GitHub Desktop.
Save m-klasen/651297e28199b4bb7907fc413c49f58f to your computer and use it in GitHub Desktop.

Get pretrained weights:

wget https://dl.fbaipublicfiles.com/detr/detr-r50-e632da11.pth

Remove class weights

checkpoint = torch.load("detr-r50-e632da11.pth", map_location='cpu')
del checkpoint["model"]["class_embed.weight"]
del checkpoint["model"]["class_embed.bias"]
torch.save(checkpoint,"detr-r50_no-class-head.pth")

and make sure to set non-strict weight loading in main.py

model_without_ddp.load_state_dict(checkpoint['model'], strict=False)

Your dataset should ideally be in the COCO-format. Make your own data-builder (alternatively rename your train/valid/annotation file to match the COCO Dataset) In datasets.coco.py add:

def build_your_dataset(image_set, args):
    root = Path(args.coco_path)
    assert root.exists(), f'provided COCO path {root} does not exist'
    mode = 'instances'
    PATHS = {
        "train": (root / "train", root / "annotations" / f'train.json'),
        "val": (root / "valid", root / "annotations" / f'valid.json'),
    }

    img_folder, ann_file = PATHS[image_set]
    dataset = CocoDetection(img_folder, ann_file, transforms=make_coco_transforms(image_set), return_masks=args.masks)
    return dataset

In datasets.__init__.py add your builder as an option:

def build_dataset(image_set, args):
    if args.dataset_file == 'coco':
        return build_coco(image_set, args)
    if args.dataset_file == 'your_dataset':
        return build_your_dataset(image_set, args)
    [...]

And lastly define how many classes you have in models.detr.py

def build(args):
    [...]
    if args.dataset_file == 'your_dataset': num_classes = 4
    [...]

Run your model (example): python main.py --dataset_file your_dataset --coco_path data --epochs 50 --lr=1e-4 --batch_size=2 --num_workers=4 --output_dir="outputs" --resume="detr-r50_no-class-head.pth"

@m-klasen
Copy link
Author

m-klasen commented Oct 8, 2020

Transformers are notoriously difficult to train, take a long time to converge (200 Epochs on COCO). Unless you can split your images into smaller subsets which feature >100 detections per crop it is going to be difficult.

@rsharmapty
Copy link

just to be clear,

  1. by no means I can change the num_queries and use the transfer learning.
  2. if I am training from scratch, I can get decent results in 200 epochs (because i see 0 mAP for 300 epochs in this case as well)

@woctezuma
Copy link

woctezuma commented Oct 8, 2020

See facebookresearch/detr#216 to work around the issue.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment