multi_gpu multi_scale

This commit is contained in:
Glenn Jocher 2019-03-19 15:43:10 +02:00
parent 2f1afd2d69
commit 735b1a370b
2 changed files with 5 additions and 1 deletions

View File

@ -105,6 +105,9 @@ 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 = \
# [], [], [], [], []
create_grids(self, 32, 1)
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
@ -225,6 +228,7 @@ def create_grids(self, img_size, nG, device='cpu'):
self.anchor_wh = self.anchor_vec.view(1, self.nA, 1, 1, 2).to(device) self.anchor_wh = self.anchor_vec.view(1, self.nA, 1, 1, 2).to(device)
self.nG = torch.FloatTensor([nG]).to(device) self.nG = torch.FloatTensor([nG]).to(device)
def load_darknet_weights(self, weights, cutoff=-1): def load_darknet_weights(self, weights, cutoff=-1):
# Parses and loads the weights stored in 'weights' # Parses and loads the weights stored in 'weights'
# cutoff: save layers between 0 and cutoff (if cutoff = -1 all are saved) # cutoff: save layers between 0 and cutoff (if cutoff = -1 all are saved)

View File

@ -12,7 +12,7 @@ def train(
data_cfg, data_cfg,
img_size=416, img_size=416,
resume=False, resume=False,
epochs=100, epochs=270,
batch_size=16, batch_size=16,
accumulate=1, accumulate=1,
multi_scale=False, multi_scale=False,