From b15ad0adb255eadffad557554cbf5af486f43521 Mon Sep 17 00:00:00 2001 From: wenzhen Date: Tue, 15 Jun 2021 07:43:01 -0700 Subject: [PATCH 1/2] Modifications for UCSD dataset --- datasets/ucsd.py | 28 ++++++++++++++++++++++++++++ main.py | 26 +++++++++++++++++++++----- models/detr.py | 8 -------- 3 files changed, 49 insertions(+), 13 deletions(-) create mode 100644 datasets/ucsd.py diff --git a/datasets/ucsd.py b/datasets/ucsd.py new file mode 100644 index 000000000..5308bfbb4 --- /dev/null +++ b/datasets/ucsd.py @@ -0,0 +1,28 @@ +""" +For UCSD dataset. + +""" +from pathlib import Path + +from .coco import CocoDetection, make_coco_transforms + +def build(image_set, args, annotation_name=None): + """ + image_set = 'train' / 'val' + category = "all/tflt/" + """ + root = Path(args.coco_path) + assert root.exists(), f'provided path {root} to custom dataset does not exist' + + if args.annotation_name is not None: + training_json_file = 'annotations_train_' + args.annotation_name + '.json' + validation_json_file = 'annotations_val_' + args.annotation_name + '.json' + + PATHS = { + "train": (root / "train2017", root / "annotations" / f'annotations_train_{annotation_name}.json'), + "val": (root / "val2017", root / "annotations" / f'annotations_val_{annotation_name}.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 \ No newline at end of file diff --git a/main.py b/main.py index e5f9eff80..e102fb0e0 100644 --- a/main.py +++ b/main.py @@ -19,12 +19,21 @@ def get_args_parser(): parser = argparse.ArgumentParser('Set transformer detector', add_help=False) + # added for UCSD + parser.add_argument('--num_classes', default=None, type=int, + help='num of classes in your dataset, which can override the value hard-coded in file models/detr.py') + parser.add_argument('--annotation_name', default=None, type=str, + help='all / tfsg / veh / tflt') + parser.add_argument('--gpu_id', default=None, type=int, + help="specify GPU ID") + + # original parser.add_argument('--lr', default=1e-4, type=float) parser.add_argument('--lr_backbone', default=1e-5, type=float) parser.add_argument('--batch_size', default=2, type=int) parser.add_argument('--weight_decay', default=1e-4, type=float) parser.add_argument('--epochs', default=300, type=int) - parser.add_argument('--lr_drop', default=200, type=int) + parser.add_argument('--lr_drop', default=100, type=int) parser.add_argument('--clip_max_norm', default=0.1, type=float, help='gradient clipping max norm') @@ -111,6 +120,7 @@ def main(args): print(args) device = torch.device(args.device) + torch.cuda.set_device(args.gpu_id) # fix the seed for reproducibility seed = args.seed + utils.get_rank() @@ -121,10 +131,14 @@ def main(args): model, criterion, postprocessors = build_model(args) model.to(device) - model_without_ddp = model if args.distributed: model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu]) model_without_ddp = model.module + else: + print("Data parallel mode activated!") + model = torch.nn.DataParallel(model, device_ids=[args.gpu_id]) + model = model.to(device) + model_without_ddp = model n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) print('number of params:', n_parameters) @@ -200,11 +214,11 @@ def main(args): if args.output_dir: checkpoint_paths = [output_dir / 'checkpoint.pth'] # extra checkpoint before LR drop and every 100 epochs - if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 100 == 0: + if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 10 == 0: checkpoint_paths.append(output_dir / f'checkpoint{epoch:04}.pth') for checkpoint_path in checkpoint_paths: utils.save_on_master({ - 'model': model_without_ddp.state_dict(), + 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'lr_scheduler': lr_scheduler.state_dict(), 'epoch': epoch, @@ -229,7 +243,7 @@ def main(args): (output_dir / 'eval').mkdir(exist_ok=True) if "bbox" in coco_evaluator.coco_eval: filenames = ['latest.pth'] - if epoch % 50 == 0: + if epoch % 10 == 0: filenames.append(f'{epoch:03}.pth') for name in filenames: torch.save(coco_evaluator.coco_eval["bbox"].eval, @@ -246,3 +260,5 @@ def main(args): if args.output_dir: Path(args.output_dir).mkdir(parents=True, exist_ok=True) main(args) + +## python -m torch.distributed.launch --nproc_per_node=8 --use_env main.py --coco_path /path/to/coco \ No newline at end of file diff --git a/models/detr.py b/models/detr.py index 23c2376da..99c1b43bf 100644 --- a/models/detr.py +++ b/models/detr.py @@ -302,14 +302,6 @@ def forward(self, x): def build(args): - # the `num_classes` naming here is somewhat misleading. - # it indeed corresponds to `max_obj_id + 1`, where max_obj_id - # is the maximum id for a class in your dataset. For example, - # COCO has a max_obj_id of 90, so we pass `num_classes` to be 91. - # As another example, for a dataset that has a single class with id 1, - # you should pass `num_classes` to be 2 (max_obj_id + 1). - # For more details on this, check the following discussion - # https://github.com/facebookresearch/detr/issues/108#issuecomment-650269223 num_classes = 20 if args.dataset_file != 'coco' else 91 if args.dataset_file == "coco_panoptic": # for panoptic, we just add a num_classes that is large enough to hold From d133196cfc09d33dfe972b58272912bca1fce507 Mon Sep 17 00:00:00 2001 From: wenzhen Date: Tue, 15 Jun 2021 08:00:45 -0700 Subject: [PATCH 2/2] Ammend the dataset --- datasets/ucsd.py | 17 ++++++++--------- main.py | 5 +++-- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/datasets/ucsd.py b/datasets/ucsd.py index 5308bfbb4..b9bcee592 100644 --- a/datasets/ucsd.py +++ b/datasets/ucsd.py @@ -6,7 +6,7 @@ from .coco import CocoDetection, make_coco_transforms -def build(image_set, args, annotation_name=None): +def build(image_set, args): """ image_set = 'train' / 'val' category = "all/tflt/" @@ -14,14 +14,13 @@ def build(image_set, args, annotation_name=None): root = Path(args.coco_path) assert root.exists(), f'provided path {root} to custom dataset does not exist' - if args.annotation_name is not None: - training_json_file = 'annotations_train_' + args.annotation_name + '.json' - validation_json_file = 'annotations_val_' + args.annotation_name + '.json' - - PATHS = { - "train": (root / "train2017", root / "annotations" / f'annotations_train_{annotation_name}.json'), - "val": (root / "val2017", root / "annotations" / f'annotations_val_{annotation_name}.json'), - } + if args.annotation_name is None: + raise ValueError("args doesn't have annotation_name") + else: + PATHS = { + "train": (root / "train2017", root / "annotations" / f'annotations_train_{args.annotation_name}.json'), + "val": (root / "val2017", root / "annotations" / f'annotations_val_{args.annotation_name}.json'), + } img_folder, ann_file = PATHS[image_set] dataset = CocoDetection(img_folder, ann_file, transforms=make_coco_transforms(image_set), return_masks=args.masks) diff --git a/main.py b/main.py index e102fb0e0..5308226eb 100644 --- a/main.py +++ b/main.py @@ -12,7 +12,8 @@ import datasets import util.misc as utils -from datasets import build_dataset, get_coco_api_from_dataset +from datasets import get_coco_api_from_dataset +from datasets.ucsd import build as build_dataset from engine import evaluate, train_one_epoch from models import build_model @@ -22,7 +23,7 @@ def get_args_parser(): # added for UCSD parser.add_argument('--num_classes', default=None, type=int, help='num of classes in your dataset, which can override the value hard-coded in file models/detr.py') - parser.add_argument('--annotation_name', default=None, type=str, + parser.add_argument('--annotation_name', default=None, type=str, action='store_true', help='all / tfsg / veh / tflt') parser.add_argument('--gpu_id', default=None, type=int, help="specify GPU ID")