From 735b1a370b177949927ee6d50644fc9125886277 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 19 Mar 2019 15:43:10 +0200 Subject: [PATCH] multi_gpu multi_scale --- models.py | 4 ++++ train.py | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/models.py b/models.py index dd77c471..18dc2b50 100755 --- a/models.py +++ b/models.py @@ -105,6 +105,9 @@ class YOLOLayer(nn.Module): self.nA = len(anchors) # number of anchors (3) self.nC = nC # number of classes (80) 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__ 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.nG = torch.FloatTensor([nG]).to(device) + def load_darknet_weights(self, weights, cutoff=-1): # Parses and loads the weights stored in 'weights' # cutoff: save layers between 0 and cutoff (if cutoff = -1 all are saved) diff --git a/train.py b/train.py index 10f6332f..5e9ee0e6 100644 --- a/train.py +++ b/train.py @@ -12,7 +12,7 @@ def train( data_cfg, img_size=416, resume=False, - epochs=100, + epochs=270, batch_size=16, accumulate=1, multi_scale=False,