multi_gpu multi_scale
This commit is contained in:
parent
bc989a0147
commit
973715060d
|
@ -105,8 +105,8 @@ class YOLOLayer(nn.Module):
|
||||||
self.nA = len(anchors) # number of anchors (3)
|
self.nA = len(anchors) # number of anchors (3)
|
||||||
self.nC = nC # number of classes (80)
|
self.nC = nC # number of classes (80)
|
||||||
self.img_size = 0
|
self.img_size = 0
|
||||||
self.nG, self.stride, self.grid_xy, self.anchor_vec, self.anchor_wh = \
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||||
[], [], [], [], []
|
create_grids(self, 32, 1, device=device)
|
||||||
|
|
||||||
if ONNX_EXPORT: # grids must be computed in __init__
|
if ONNX_EXPORT: # grids must be computed in __init__
|
||||||
stride = [32, 16, 8][yolo_layer] # stride of this layer
|
stride = [32, 16, 8][yolo_layer] # stride of this layer
|
||||||
|
|
Loading…
Reference in New Issue