glenn-jocher commited on
Commit
2703ac7
1 Parent(s): 4e32b60

profiling update

Browse files
Files changed (1) hide show
  1. models/yolo.py +7 -7
models/yolo.py CHANGED
@@ -20,6 +20,7 @@ class Detect(nn.Module):
20
  self.export = False # onnx export
21
 
22
  def forward(self, x):
 
23
  z = [] # inference output
24
  self.training |= self.export
25
  for i in range(self.nl):
@@ -66,7 +67,7 @@ class Model(nn.Module):
66
  print('')
67
 
68
  def forward(self, x, augment=False, profile=False):
69
- y, ts = [], 0 # outputs
70
  for m in self.model:
71
  if m.f != -1: # if not from previous layer
72
  x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
@@ -77,15 +78,14 @@ class Model(nn.Module):
77
  t = torch_utils.time_synchronized()
78
  for _ in range(10):
79
  _ = m(x)
80
- dt = torch_utils.time_synchronized() - t
81
- ts += dt
82
- print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt * 100, m.type))
83
 
84
  x = m(x) # run
85
  y.append(x if m.i in self.save else None) # save output
86
 
87
  if profile:
88
- print(ts * 100)
89
  return x
90
 
91
  def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency
@@ -180,8 +180,8 @@ if __name__ == '__main__':
180
  model.train()
181
 
182
  # Profile
183
- # img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
184
- # y = model(img, profile=True)
185
  # print([y[0].shape] + [x.shape for x in y[1]])
186
 
187
  # ONNX export
 
20
  self.export = False # onnx export
21
 
22
  def forward(self, x):
23
+ x = x.copy() # for profiling
24
  z = [] # inference output
25
  self.training |= self.export
26
  for i in range(self.nl):
 
67
  print('')
68
 
69
  def forward(self, x, augment=False, profile=False):
70
+ y, dt = [], [] # outputs
71
  for m in self.model:
72
  if m.f != -1: # if not from previous layer
73
  x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
 
78
  t = torch_utils.time_synchronized()
79
  for _ in range(10):
80
  _ = m(x)
81
+ dt.append((torch_utils.time_synchronized() - t) * 100)
82
+ print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type))
 
83
 
84
  x = m(x) # run
85
  y.append(x if m.i in self.save else None) # save output
86
 
87
  if profile:
88
+ print('%.1fms total' % sum(dt))
89
  return x
90
 
91
  def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency
 
180
  model.train()
181
 
182
  # Profile
183
+ img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
184
+ y = model(img, profile=True)
185
  # print([y[0].shape] + [x.shape for x in y[1]])
186
 
187
  # ONNX export