Yuantao Yang yangyuantao commited on
Commit
e7fc274
1 Parent(s): ce7deec

Fix module count in parse_model (#4379)

Browse files

Co-authored-by: yangyuantao <[email protected]>

Files changed (1) hide show
  1. models/yolo.py +2 -2
models/yolo.py CHANGED
@@ -234,7 +234,7 @@ def parse_model(d, ch): # model_dict, input_channels(3)
234
  except:
235
  pass
236
 
237
- n = max(round(n * gd), 1) if n > 1 else n # depth gain
238
  if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP,
239
  C3, C3TR, C3SPP]:
240
  c1, c2 = ch[f], args[0]
@@ -264,7 +264,7 @@ def parse_model(d, ch): # model_dict, input_channels(3)
264
  t = str(m)[8:-2].replace('__main__.', '') # module type
265
  np = sum([x.numel() for x in m_.parameters()]) # number params
266
  m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params
267
- LOGGER.info('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print
268
  save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
269
  layers.append(m_)
270
  if i == 0:
 
234
  except:
235
  pass
236
 
237
+ n = n_ = max(round(n * gd), 1) if n > 1 else n # depth gain
238
  if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP,
239
  C3, C3TR, C3SPP]:
240
  c1, c2 = ch[f], args[0]
 
264
  t = str(m)[8:-2].replace('__main__.', '') # module type
265
  np = sum([x.numel() for x in m_.parameters()]) # number params
266
  m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params
267
+ LOGGER.info('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n_, np, t, args)) # print
268
  save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
269
  layers.append(m_)
270
  if i == 0: