profiling update

pull/1/head
Glenn Jocher 5 years ago
parent 4e32b6029a
commit 2703ac73e8

@ -20,6 +20,7 @@ class Detect(nn.Module):
self.export = False # onnx export self.export = False # onnx export
def forward(self, x): def forward(self, x):
x = x.copy() # for profiling
z = [] # inference output z = [] # inference output
self.training |= self.export self.training |= self.export
for i in range(self.nl): for i in range(self.nl):
@ -66,7 +67,7 @@ class Model(nn.Module):
print('') print('')
def forward(self, x, augment=False, profile=False): def forward(self, x, augment=False, profile=False):
y, ts = [], 0 # outputs y, dt = [], [] # outputs
for m in self.model: for m in self.model:
if m.f != -1: # if not from previous layer if m.f != -1: # if not from previous layer
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 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):
t = torch_utils.time_synchronized() t = torch_utils.time_synchronized()
for _ in range(10): for _ in range(10):
_ = m(x) _ = m(x)
dt = torch_utils.time_synchronized() - t dt.append((torch_utils.time_synchronized() - t) * 100)
ts += dt print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type))
print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt * 100, m.type))
x = m(x) # run x = m(x) # run
y.append(x if m.i in self.save else None) # save output y.append(x if m.i in self.save else None) # save output
if profile: if profile:
print(ts * 100) print('%.1fms total' % sum(dt))
return x return x
def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency
@ -180,8 +180,8 @@ if __name__ == '__main__':
model.train() model.train()
# Profile # Profile
# img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device) img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
# y = model(img, profile=True) y = model(img, profile=True)
# print([y[0].shape] + [x.shape for x in y[1]]) # print([y[0].shape] + [x.shape for x in y[1]])
# ONNX export # ONNX export

Loading…
Cancel
Save