This commit is contained in:
Glenn Jocher 2020-02-17 12:36:11 -08:00
parent aa45dc05b3
commit 426d5b82c6
1 changed files with 3 additions and 3 deletions

View File

@ -54,10 +54,10 @@ def create_modules(module_defs, img_size, arc):
elif mdef['type'] == 'upsample': elif mdef['type'] == 'upsample':
if ONNX_EXPORT: # explicitly state size, avoid scale_factor if ONNX_EXPORT: # explicitly state size, avoid scale_factor
g = (yolo_index + 1) * 2 # gain g = (yolo_index + 1) * 2 / 32 # gain
modules = nn.Upsample(size=(10 * g, 6 * g), mode='nearest') # assume img_size = (320, 192) modules = nn.Upsample(size=tuple(int(x * g) for x in img_size)) # img_size = (320, 192)
else: else:
modules = nn.Upsample(scale_factor=int(mdef['stride']), mode='nearest') modules = nn.Upsample(scale_factor=int(mdef['stride']))
elif mdef['type'] == 'route': # nn.Sequential() placeholder for 'route' layer elif mdef['type'] == 'route': # nn.Sequential() placeholder for 'route' layer
layers = [int(x) for x in mdef['layers'].split(',')] layers = [int(x) for x in mdef['layers'].split(',')]