def benchmark_eval(args):
cfg = setup(args)
model = build_model(cfg)
model.eval()
logger.info("Model:\n{}".format(model))
DetectionCheckpointer(model).load(cfg.MODEL.WEIGHTS)
cfg.defrost()
cfg.DATALOADER.NUM_WORKERS = 0
data_loader = build_detection_test_loader(cfg, cfg.DATASETS.TEST[0])
dummy_data = list(itertools.islice(data_loader, 100))
def f():
while True:
yield from DatasetFromList(dummy_data, copy=False)
for _ in range(5): # warmup
model(dummy_data[0])
max_iter = 400
timer = Timer()
with tqdm.tqdm(total=max_iter) as pbar:
for idx, d in enumerate(f()):
if idx == max_iter:
break
model(d)
pbar.update()
logger.info("{} iters in {} seconds.".format(max_iter, timer.seconds()))
def benchmark_train(args):
cfg = setup(args)
model = build_model(cfg)
logger.info("Model:\n{}".format(model))
if comm.get_world_size() > 1:
model = DistributedDataParallel(model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False)
optimizer = build_optimizer(cfg, model)
checkpointer = DetectionCheckpointer(model, optimizer=optimizer)
checkpointer.load(cfg.MODEL.WEIGHTS)
cfg.defrost()
cfg.DATALOADER.NUM_WORKERS = 0
data_loader = build_detection_train_loader(cfg)
dummy_data = list(itertools.islice(data_loader, 100))
def f():
while True:
yield from DatasetFromList(dummy_data, copy=False)
max_iter = 400
trainer = SimpleTrainer(model, f(), optimizer)
trainer.register_hooks([
hooks.IterationTimer(),
hooks.PeriodicWriter([CommonMetricPrinter(max_iter)])
])
trainer.train(1, max_iter)
parser.add_argument("--output",
help="output directory for the converted caffe2 model")
parser.add_argument(
"opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
logger = setup_logger()
logger.info("Command line arguments: " + str(args))
cfg = setup_cfg(args)
# create a torch model
torch_model = build_model(cfg)
DefaultCheckpointer(torch_model).resume_or_load(cfg.MODEL.WEIGHTS)
# get a sample data
data_loader = build_test_loader(cfg, cfg.DATASETS.TEST[0])
first_batch = next(iter(data_loader))
# convert and save caffe2 model
caffe2_model = export_caffe2_model(cfg, torch_model, first_batch)
caffe2_model.save_protobuf(args.output)
# draw the caffe2 graph
caffe2_model.save_graph(os.path.join(args.output, "model.svg"),
inputs=first_batch)
# run evaluation with the converted model
if args.run_eval: