cleanup/speedup
This commit is contained in:
parent
1b68fe7fde
commit
fff45c39a8
|
@ -105,7 +105,10 @@ def create_modules(module_defs, img_size):
|
||||||
module_list.append(modules)
|
module_list.append(modules)
|
||||||
output_filters.append(filters)
|
output_filters.append(filters)
|
||||||
|
|
||||||
return module_list, routs
|
routs_binary = [False] * (i + 1)
|
||||||
|
for i in routs:
|
||||||
|
routs_binary[i] = True
|
||||||
|
return module_list, routs_binary
|
||||||
|
|
||||||
|
|
||||||
class weightedFeatureFusion(nn.Module): # weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070
|
class weightedFeatureFusion(nn.Module): # weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070
|
||||||
|
@ -292,7 +295,7 @@ class Darknet(nn.Module):
|
||||||
# print(''), [print(out[i].shape) for i in layers], print(x.shape)
|
# print(''), [print(out[i].shape) for i in layers], print(x.shape)
|
||||||
elif mtype == 'yolo':
|
elif mtype == 'yolo':
|
||||||
yolo_out.append(module(x, img_size, out))
|
yolo_out.append(module(x, img_size, out))
|
||||||
out.append(x if i in self.routs else [])
|
out.append(x if self.routs[i] else [])
|
||||||
if verbose:
|
if verbose:
|
||||||
print('%g/%g %s -' % (i, len(self.module_list), mtype), list(x.shape), str)
|
print('%g/%g %s -' % (i, len(self.module_list), mtype), list(x.shape), str)
|
||||||
str = ''
|
str = ''
|
||||||
|
@ -342,7 +345,7 @@ def create_grids(self, img_size=416, ng=(13, 13), device='cpu', type=torch.float
|
||||||
|
|
||||||
# build wh gains
|
# build wh gains
|
||||||
self.anchor_vec = self.anchors.to(device) / self.stride
|
self.anchor_vec = self.anchors.to(device) / self.stride
|
||||||
self.anchor_wh = self.anchor_vec.view(1, self.na, 1, 1, 2).to(device).type(type)
|
self.anchor_wh = self.anchor_vec.view(1, self.na, 1, 1, 2).type(type)
|
||||||
self.ng = torch.Tensor(ng).to(device)
|
self.ng = torch.Tensor(ng).to(device)
|
||||||
self.nx = nx
|
self.nx = nx
|
||||||
self.ny = ny
|
self.ny = ny
|
||||||
|
|
Loading…
Reference in New Issue