|
|
@ -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
|
|
|
|