glenn-jocher
commited on
Commit
•
2703ac7
1
Parent(s):
4e32b60
profiling update
Browse files- 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,
|
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
|
81 |
-
|
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(
|
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 |
-
|
184 |
-
|
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
|