Add efficientdet model

Author: xuzhao9

torchbenchmark/models/timm_efficientdet/__init__.py

@@ -1,14 +1,10 @@
-import torch
 import os
-import random
 import logging
-from contextlib import suppress
+import torch
 from pathlib import Path
 
 # TorchBench imports
 from torchbenchmark.util.model import BenchmarkModel
-from torchbenchmark.util.jit import jit_model
-from torchbenchmark.util.torch_feature_checker import check_native_amp
 from torchbenchmark.tasks import COMPUTER_VISION
 
 # effdet imports
@@ -18,7 +14,7 @@
 # timm imports
 from timm.models.layers import set_layer_config
 from timm.optim import create_optimizer
-from timm.utils import ModelEmaV2, NativeScaler
+from timm.utils import ModelEmaV2
 from timm.scheduler import create_scheduler
 
 # local imports
@@ -30,129 +26,107 @@
 CURRENT_DIR = Path(os.path.dirname(os.path.realpath(__file__)))
 DATA_DIR = os.path.join(CURRENT_DIR.parent.parent, "data", ".data", "coco2017-minimal", "coco")
 
-torch.manual_seed(1337)
-random.seed(1337)
-torch.backends.cudnn.deterministic = False
-torch.backends.cudnn.benchmark = True
-
 class Model(BenchmarkModel):
     task = COMPUTER_VISION.DETECTION
-
     # Original Train batch size 32 on 2x RTX 3090 (24 GB cards)
     # Downscale to batch size 16 on single GPU
-    def __init__(self, device=None, jit=False, train_bs=16, eval_bs=128):
-        super().__init__()
-        self.device = device
-        self.jit = jit
+    DEFAULT_TRAIN_BSIZE = 16
+    DEFAULT_EVAL_BSIZE = 128
+
+    def __init__(self, test, device, jit=False, batch_size=None, extra_args=[]):
+        super().__init__(test=test, device=device, jit=jit, batch_size=batch_size, extra_args=extra_args)
         # generate arguments
         args = get_args()
         # setup train and eval batch size
-        args.batch_size = train_bs
-        args.eval_batch_size = eval_bs
-        # Use native amp if possible
-        args.native_amp = check_native_amp()
+        args.batch_size = self.batch_size
         # Disable distributed
         args.distributed = False
-        args.device = device
-        args.torchscript = jit
+        args.device = self.device
+        args.torchscript = self.jit
         args.world_size = 1
         args.rank = 0
         args.pretrained_backbone = not args.no_pretrained_backbone
         args.prefetcher = not args.no_prefetcher
         args.root = DATA_DIR
 
-        if not self.device == "cuda":
-            raise NotImplementedError("Only CUDA is supported by this model") 
-
         with set_layer_config(scriptable=args.torchscript):
-            extra_args = {}
+            timm_extra_args = {}
             if args.img_size is not None:
-                extra_args = dict(image_size=(args.img_size, args.img_size))
-            model = create_model(
-                model_name=args.model,
-                bench_task='train',
-                num_classes=args.num_classes,
-                pretrained=args.pretrained,
-                pretrained_backbone=args.pretrained_backbone,
-                redundant_bias=args.redundant_bias,
-                label_smoothing=args.smoothing,
-                legacy_focal=args.legacy_focal,
-                jit_loss=args.jit_loss,
-                soft_nms=args.soft_nms,
-                bench_labeler=args.bench_labeler,
-                checkpoint_path=args.initial_checkpoint,
-            )
-            eval_model = create_model(
-                model_name=args.model,
-                bench_task='predict',
-                num_classes=args.num_classes,
-                pretrained=args.pretrained,
-                redundant_bias=args.redundant_bias,
-                soft_nms=args.soft_nms,
-                checkpoint_path=args.checkpoint,
-                checkpoint_ema=args.use_ema,
-                **extra_args,
-            )
+                timm_extra_args = dict(image_size=(args.img_size, args.img_size))
+            if test == "train":
+                model = create_model(
+                    model_name=args.model,
+                    bench_task='train',
+                    num_classes=args.num_classes,
+                    pretrained=args.pretrained,
+                    pretrained_backbone=args.pretrained_backbone,
+                    redundant_bias=args.redundant_bias,
+                    label_smoothing=args.smoothing,
+                    legacy_focal=args.legacy_focal,
+                    jit_loss=args.jit_loss,
+                    soft_nms=args.soft_nms,
+                    bench_labeler=args.bench_labeler,
+                    checkpoint_path=args.initial_checkpoint,
+                )
+            elif test == "eval":
+                model = create_model(
+                    model_name=args.model,
+                    bench_task='predict',
+                    num_classes=args.num_classes,
+                    pretrained=args.pretrained,
+                    redundant_bias=args.redundant_bias,
+                    soft_nms=args.soft_nms,
+                    checkpoint_path=args.checkpoint,
+                    checkpoint_ema=args.use_ema,
+                    **timm_extra_args,
+                )
         model_config = model.config  # grab before we obscure with DP/DDP wrappers
-        model = model.to(device)
+        self.model = model.to(device)
         if args.channels_last:
-            model = model.to(memory_format=torch.channels_last)
-        eval_model = eval_model.to(device)
-
-        self.model, self.eval_model = jit_model(model, eval_model, jit=jit)
-        self.optimizer = create_optimizer(args, model)
-        self.amp_autocast = suppress
-        if args.native_amp:
-            self.amp_autocast = torch.cuda.amp.autocast
-            self.loss_scaler = NativeScaler()
-        self.model_ema = None
-        if args.model_ema:
-            # Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
-            self.model_ema = ModelEmaV2(model, decay=args.model_ema_decay)
-        self.lr_scheduler, self.num_epochs = create_scheduler(args, self.optimizer)
-
-        self.loader_train, self.loader_eval, self.evaluator, _, dataset_eval = create_datasets_and_loaders(args, model_config)
-        if model_config.num_classes < self.loader_train.dataset.parser.max_label:
-            logging.error(
-                f'Model {model_config.num_classes} has fewer classes than dataset {self.loader_train.dataset.parser.max_label}.')
-            exit(1)
-        if model_config.num_classes > self.loader_train.dataset.parser.max_label:
-            logging.warning(
-                f'Model {model_config.num_classes} has more classes than dataset {self.loader_train.dataset.parser.max_label}.')
-        self.train_num_batch = 1
-
-        # Create eval loader
-        input_config = resolve_input_config(args, model_config)
-        self.eval_loader = create_loader(
-                dataset_eval,
-                input_size=input_config['input_size'],
-                batch_size=args.eval_batch_size,
-                use_prefetcher=args.prefetcher,
-                interpolation=args.eval_interpolation,
-                fill_color=input_config['fill_color'],
-                mean=input_config['mean'],
-                std=input_config['std'],
-                num_workers=args.workers,
-                pin_mem=args.pin_mem)
-        self.eval_num_batch = 1
+            self.model = self.model.to(memory_format=torch.channels_last)
+
+        if test == "train":
+            self.optimizer = create_optimizer(args, model)
+            self.model_ema = None
+            if args.model_ema:
+                # Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
+                self.model_ema = ModelEmaV2(model, decay=args.model_ema_decay)
+            self.lr_scheduler, self.num_epochs = create_scheduler(args, self.optimizer)
+
+            self.loader_train, self.loader_eval, self.evaluator, _, dataset_eval = create_datasets_and_loaders(args, model_config)
+            if model_config.num_classes < self.loader_train.dataset.parser.max_label:
+                logging.error(
+                    f'Model {model_config.num_classes} has fewer classes than dataset {self.loader_train.dataset.parser.max_label}.')
+                exit(1)
+            if model_config.num_classes > self.loader_train.dataset.parser.max_label:
+                logging.warning(
+                    f'Model {model_config.num_classes} has more classes than dataset {self.loader_train.dataset.parser.max_label}.')
+        elif test == "eval":
+            # Create eval loader
+            input_config = resolve_input_config(args, model_config)
+            self.loader = create_loader(
+                    dataset_eval,
+                    input_size=input_config['input_size'],
+                    batch_size=args.eval_batch_size,
+                    use_prefetcher=args.prefetcher,
+                    interpolation=args.eval_interpolation,
+                    fill_color=input_config['fill_color'],
+                    mean=input_config['mean'],
+                    std=input_config['std'],
+                    num_workers=args.workers,
+                    pin_mem=args.pin_mem)
         self.args = args
+        # Only run 1 batch in 1 epoch
+        self.num_batches = 1
+        self.num_epochs = 1
 
     def get_module(self):
-        self.eval_model.eval()
-        for _, (input, target) in zip(range(self.eval_num_batch), self.loader_eval):
-            return (self.eval_model, (input, target))
+        for _, (input, target) in zip(range(self.num_batches), self.loader_eval):
+            return (self.model, (input, target))
 
-    # Temporarily disable training because this will cause CUDA OOM in CI
-    # TODO: re-enable this test when better hardware is available
     def train(self, niter=1):
-        raise NotImplementedError("Disable this test because it causes CUDA OOM on Nvidia T4 GPU")
-        if not self.device == "cuda":
-            raise NotImplementedError("Only CUDA is supported by this model")
-        if self.jit:
-            raise NotImplementedError("JIT is not supported by this model")
         eval_metric = self.args.eval_metric
-        self.model.train()
-        for epoch in range(niter):
+        for epoch in range(self.num_epochs):
             train_metrics = train_epoch(
                 epoch, self.model, self.loader_train,
                 self.optimizer, self.args,
@@ -170,14 +144,8 @@ def train(self, niter=1):
                 self.lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
 
     def eval(self, niter=1):
-        if not self.device == "cuda":
-            raise NotImplementedError("Only CUDA is supported by this model")
-        if self.jit:
-            raise NotImplementedError("JIT is not supported by this model")
-        self.eval_model.eval()
         for _ in range(niter):
             with torch.no_grad():
-                for _, (input, target) in zip(range(self.eval_num_batch), self.eval_loader):
-                    with self.amp_autocast():
-                        output = self.eval_model(input, img_info=target)
+                for _, (input, target) in zip(range(self.num_batches), self.loader):
+                    output = self.model(input, img_info=target)
                     self.evaluator.add_predictions(output, target)

torchbenchmark/models/timm_efficientdet/args.py

@@ -170,8 +170,6 @@ def _parse_args():
     parser.add_argument("--local_rank", default=0, type=int)
 
     # Evaluation parameters
-    parser.add_argument('--eval-batch-size', default=128, type=int,
-                metavar='N', help='eval mini-batch size (default: 128)')
     parser.add_argument('--eval-interpolation', default='bilinear', type=str, metavar='NAME',
                 help='Image resize interpolation type (overrides model)')
     parser.add_argument('--img-size', default=None, type=int,

torchbenchmark/models/timm_efficientdet/metadata.yaml

@@ -0,0 +1,11 @@
+eval_benchmark: false
+eval_deterministic: false
+eval_nograd: true
+optimized_for_inference: false
+train_benchmark: false
+train_deterministic: false
+not_implemented:
+  # Disable CUDA test because of insufficient GPU memory on CI machine
+  - device: cuda
+  # CPU is not supported
+  - device: cpu