glenn-jocher
commited on
Commit
•
c672bef
1
Parent(s):
12b0c04
model fuse
Browse files
detect.py
CHANGED
@@ -21,6 +21,8 @@ def detect(save_img=False):
|
|
21 |
google_utils.attempt_download(weights)
|
22 |
model = torch.load(weights, map_location=device)['model']
|
23 |
# torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning
|
|
|
|
|
24 |
|
25 |
# Second-stage classifier
|
26 |
classify = False
|
@@ -29,12 +31,6 @@ def detect(save_img=False):
|
|
29 |
modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']) # load weights
|
30 |
modelc.to(device).eval()
|
31 |
|
32 |
-
# Eval mode
|
33 |
-
model.to(device).eval()
|
34 |
-
|
35 |
-
# Fuse Conv2d + BatchNorm2d layers
|
36 |
-
# model.fuse()
|
37 |
-
|
38 |
# Half precision
|
39 |
half = half and device.type != 'cpu' # half precision only supported on CUDA
|
40 |
if half:
|
|
|
21 |
google_utils.attempt_download(weights)
|
22 |
model = torch.load(weights, map_location=device)['model']
|
23 |
# torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning
|
24 |
+
# model.fuse()
|
25 |
+
model.to(device).eval()
|
26 |
|
27 |
# Second-stage classifier
|
28 |
classify = False
|
|
|
31 |
modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']) # load weights
|
32 |
modelc.to(device).eval()
|
33 |
|
|
|
|
|
|
|
|
|
|
|
|
|
34 |
# Half precision
|
35 |
half = half and device.type != 'cpu' # half precision only supported on CUDA
|
36 |
if half:
|