Skip to content

Commit

Permalink
Merge pull request #434 from developer0hye/patch-3
Browse files Browse the repository at this point in the history
cover do_detect function with torch.no_grad():
  • Loading branch information
Tianxiaomo committed Dec 12, 2021
2 parents 0f7754a + 3321f55 commit 47483be
Showing 1 changed file with 21 additions and 20 deletions.
41 changes: 21 additions & 20 deletions tool/torch_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,30 +75,31 @@ def convert2cpu_long(gpu_matrix):

def do_detect(model, img, conf_thresh, nms_thresh, use_cuda=1):
model.eval()
t0 = time.time()
with torch.no_grad():
t0 = time.time()

if type(img) == np.ndarray and len(img.shape) == 3: # cv2 image
img = torch.from_numpy(img.transpose(2, 0, 1)).float().div(255.0).unsqueeze(0)
elif type(img) == np.ndarray and len(img.shape) == 4:
img = torch.from_numpy(img.transpose(0, 3, 1, 2)).float().div(255.0)
else:
print("unknow image type")
exit(-1)
if type(img) == np.ndarray and len(img.shape) == 3: # cv2 image
img = torch.from_numpy(img.transpose(2, 0, 1)).float().div(255.0).unsqueeze(0)
elif type(img) == np.ndarray and len(img.shape) == 4:
img = torch.from_numpy(img.transpose(0, 3, 1, 2)).float().div(255.0)
else:
print("unknow image type")
exit(-1)

if use_cuda:
img = img.cuda()
img = torch.autograd.Variable(img)

if use_cuda:
img = img.cuda()
img = torch.autograd.Variable(img)

t1 = time.time()
t1 = time.time()

output = model(img)
output = model(img)

t2 = time.time()
t2 = time.time()

print('-----------------------------------')
print(' Preprocess : %f' % (t1 - t0))
print(' Model Inference : %f' % (t2 - t1))
print('-----------------------------------')
print('-----------------------------------')
print(' Preprocess : %f' % (t1 - t0))
print(' Model Inference : %f' % (t2 - t1))
print('-----------------------------------')

return utils.post_processing(img, conf_thresh, nms_thresh, output)
return utils.post_processing(img, conf_thresh, nms_thresh, output)

0 comments on commit 47483be

Please sign in to comment.