add NMS to pretrained pytorch hub models

This commit is contained in:
Glenn Jocher committed 2020-09-18 18:17:11 -07:00
1 parent 5a9c5c1d3b
commit c4cb78570c
2 files changed
+21

No files matched your search

+7
View File
@@ -10,6 +10,7 @@ import os
import torch
from models.common import NMS
from models.yolo import Model
from utils.google_utils import attempt_download
@@ -35,6 +36,12 @@ def create(name, pretrained, channels, classes):
state_dict = torch.load(ckpt, map_location=torch.device('cpu'))['model'].float().state_dict() # to FP32
state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape} # filter
model.load_state_dict(state_dict, strict=False) # load
m = NMS()
m.f = -1 # from
m.i = model.model[-1].i + 1 # index
model.model.add_module(name='%s' % m.i, module=m) # add NMS
model.eval()
return model
except Exception as e: