From 8699c319ca48253eb9983ee3be4c7ab0c3159d60 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 24 Jun 2020 10:42:01 -0700 Subject: [PATCH] update yolo.py --- models/yolo.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/models/yolo.py b/models/yolo.py index 158a241..2d853cb 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -96,8 +96,11 @@ class Model(nn.Module): x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers if profile: - import thop - o = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # FLOPS + try: + import thop + o = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # FLOPS + except: + o = 0 t = torch_utils.time_synchronized() for _ in range(10): _ = m(x) @@ -217,11 +220,10 @@ if __name__ == '__main__': # Profile # img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device) # y = model(img, profile=True) - # print([y[0].shape] + [x.shape for x in y[1]]) # ONNX export # model.model[-1].export = True - # torch.onnx.export(model, img, f.replace('.yaml', '.onnx'), verbose=True, opset_version=11) + # torch.onnx.export(model, img, opt.cfg.replace('.yaml', '.onnx'), verbose=True, opset_version=11) # Tensorboard # from torch.utils.tensorboard import SummaryWriter