in torchbenchmark/models/timm_efficientdet/__init__.py [0:0]
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)
if not device == "cuda":
# Only implemented on CUDA because the original model code explicitly calls the `Tensor.cuda()` API
# https://github.com/rwightman/efficientdet-pytorch/blob/9cb43186711d28bd41f82f132818c65663b33c1f/effdet/data/loader.py#L114
raise NotImplementedError("The original model code forces the use of CUDA.")
# generate arguments
args = get_args()
# setup train and eval batch size
args.batch_size = self.batch_size
# Disable distributed
args.distributed = False
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
with set_layer_config(scriptable=args.torchscript):
timm_extra_args = {}
if args.img_size is not None:
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
self.model = model.to(device)
if args.channels_last:
self.model = self.model.to(memory_format=torch.channels_last)
self.loader_train, self.loader_eval, self.evaluator, _, dataset_eval = create_datasets_and_loaders(args, model_config)
self.amp_autocast = suppress
if test == "train":
self.optimizer = create_optimizer(args, model)
self.loss_scaler = None
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)
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.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