Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Fuse by default when test.py called directly (faster)

This commit is contained in:
Glenn Jocher 2020-03-19 18:11:08 -07:00
parent fff45c39a8
commit 89b6377723
2 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
models.py
View File

@ -311,7 +311,7 @@ class Darknet(nn.Module):
def fuse(self): def fuse(self):
# Fuse Conv2d + BatchNorm2d layers throughout model # Fuse Conv2d + BatchNorm2d layers throughout model
print('Fusing Conv2d() and BatchNorm2d() layers...') print('Fusing layers...')
fused_list = nn.ModuleList() fused_list = nn.ModuleList()
for a in list(self.children())[0]: for a in list(self.children())[0]:
if isinstance(a, nn.Sequential): if isinstance(a, nn.Sequential):

6
test.py
View File

@ -29,7 +29,7 @@ def test(cfg,
os.remove(f) os.remove(f)
# Initialize model # Initialize model
model = Darknet(cfg, img_size).to(device) model = Darknet(cfg, img_size)
# Load weights # Load weights
attempt_download(weights) attempt_download(weights)
@ -38,6 +38,10 @@ def test(cfg,
else: # darknet format else: # darknet format
load_darknet_weights(model, weights) load_darknet_weights(model, weights)
# Fuse
model.fuse()
model.to(device)
if device.type != 'cpu' and torch.cuda.device_count() > 1: if device.type != 'cpu' and torch.cuda.device_count() > 1:
model = nn.DataParallel(model) model = nn.DataParallel(model)
else: # called by train.py else: # called by train.py