Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

multi_gpu multi_scale

This commit is contained in:
Glenn Jocher 2019-03-19 10:38:32 +02:00
parent 32f1def48f
commit 056eed2833
5 changed files with 32 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
detect.py
View File

@ -36,7 +36,7 @@ def detect(
os.system('wget https://storage.googleapis.com/ultralytics/yolov3.pt -O ' + weights) os.system('wget https://storage.googleapis.com/ultralytics/yolov3.pt -O ' + weights)
model.load_state_dict(torch.load(weights, map_location='cpu')['model']) model.load_state_dict(torch.load(weights, map_location='cpu')['model'])
else: # darknet format else: # darknet format
load_darknet_weights(model, weights) _ = load_darknet_weights(model, weights)
model.to(device).eval() model.to(device).eval()

6
models.py
View File

@ -293,11 +293,7 @@ def load_darknet_weights(self, weights, cutoff=-1):
conv_layer.weight.data.copy_(conv_w) conv_layer.weight.data.copy_(conv_w)
ptr += num_w ptr += num_w
return cutoff
"""
@:param path - path of the new weights file
@:param cutoff - save layers between 0 and cutoff (cutoff = -1 -> all are saved)
"""
def save_weights(self, path, cutoff=-1): def save_weights(self, path, cutoff=-1):

2
test.py
View File

@ -35,7 +35,7 @@ def test(
if weights.endswith('.pt'): # pytorch format if weights.endswith('.pt'): # pytorch format
model.load_state_dict(torch.load(weights, map_location='cpu')['model']) model.load_state_dict(torch.load(weights, map_location='cpu')['model'])
else: # darknet format else: # darknet format
load_darknet_weights(model, weights) _ = load_darknet_weights(model, weights)
model.to(device).eval() model.to(device).eval()

47
train.py
View File

@ -14,7 +14,7 @@ def train(
resume=False, resume=False,
epochs=100, epochs=100,
batch_size=16, batch_size=16,
accumulated_batches=1, accumulate=1,
multi_scale=False, multi_scale=False,
freeze_backbone=False, freeze_backbone=False,
): ):
@ -35,9 +35,9 @@ def train(
model = Darknet(cfg, img_size) model = Darknet(cfg, img_size)
# Get dataloader # Get dataloader
dataloader = LoadImagesAndLabels(train_path, batch_size, img_size, multi_scale=multi_scale, augment=True) dataloader = LoadImagesAndLabels(train_path, batch_size, img_size, augment=True)
lr0 = 0.001 lr0 = 0.001 # initial learning rate
cutoff = -1 # backbone reaches to cutoff layer cutoff = -1 # backbone reaches to cutoff layer
start_epoch = 0 start_epoch = 0
best_loss = float('inf') best_loss = float('inf')
@ -64,14 +64,12 @@ def train(
else: else:
# Initialize model with backbone (optional) # Initialize model with backbone (optional)
if cfg.endswith('yolov3.cfg'): if cfg.endswith('yolov3.cfg'):
load_darknet_weights(model, weights + 'darknet53.conv.74') cutoff = load_darknet_weights(model, weights + 'darknet53.conv.74')
cutoff = 75
elif cfg.endswith('yolov3-tiny.cfg'): elif cfg.endswith('yolov3-tiny.cfg'):
load_darknet_weights(model, weights + 'yolov3-tiny.conv.15') cutoff = load_darknet_weights(model, weights + 'yolov3-tiny.conv.15')
cutoff = 15
# Set optimizer # Set optimizer
optimizer = torch.optim.SGD(filter(lambda x: x.requires_grad, model.parameters()), lr=lr0, momentum=.9) optimizer = torch.optim.SGD(model.parameters(), lr=lr0, momentum=.9)
if torch.cuda.device_count() > 1: if torch.cuda.device_count() > 1:
model = nn.DataParallel(model) model = nn.DataParallel(model)
@ -94,22 +92,21 @@ def train(
# Update scheduler (automatic) # Update scheduler (automatic)
# scheduler.step() # scheduler.step()
# Update scheduler (manual) at 0, 54, 61 epochs to 1e-3, 1e-4, 1e-5 # Update scheduler (manual)
if epoch > 250: if epoch > 250:
lr = lr0 / 10 lr = lr0 / 10
else: else:
lr = lr0 lr = lr0
for g in optimizer.param_groups: for x in optimizer.param_groups:
g['lr'] = lr x['lr'] = lr
# Freeze darknet53.conv.74 for first epoch # Freeze backbone at epoch 0, unfreeze at epoch 1
if freeze_backbone and (epoch < 2): if freeze_backbone and epoch < 2:
for i, (name, p) in enumerate(model.named_parameters()): for i, (name, p) in enumerate(model.named_parameters()):
if int(name.split('.')[1]) < cutoff: # if layer < 75 if int(name.split('.')[1]) < cutoff: # if layer < 75
p.requires_grad = False if (epoch == 0) else True p.requires_grad = False if (epoch == 0) else True
ui = -1 ui = -1
optimizer.zero_grad()
rloss = defaultdict(float) rloss = defaultdict(float)
for i, (imgs, targets, _, _) in enumerate(dataloader): for i, (imgs, targets, _, _) in enumerate(dataloader):
targets = targets.to(device) targets = targets.to(device)
@ -118,10 +115,10 @@ def train(
continue continue
# SGD burn-in # SGD burn-in
if (epoch == 0) & (i <= n_burnin): if (epoch == 0) and (i <= n_burnin):
lr = lr0 * (i / n_burnin) ** 4 lr = lr0 * (i / n_burnin) ** 4
for g in optimizer.param_groups: for x in optimizer.param_groups:
g['lr'] = lr x['lr'] = lr
# Run model # Run model
pred = model(imgs.to(device)) pred = model(imgs.to(device))
@ -136,7 +133,7 @@ def train(
loss.backward() loss.backward()
# Accumulate gradient for x batches before optimizing # Accumulate gradient for x batches before optimizing
if ((i + 1) % accumulated_batches == 0) or (i == len(dataloader) - 1): if (i + 1) % accumulate == 0 or (i + 1) == len(dataloader):
optimizer.step() optimizer.step()
optimizer.zero_grad() optimizer.zero_grad()
@ -154,11 +151,17 @@ def train(
t0 = time.time() t0 = time.time()
print(s) print(s)
# Multi-Scale training (320 - 608 pixels) every 10 batches
if multi_scale and (i + 1) % 10 == 0:
dataloader.img_size = random.choice(range(10, 20)) * 32
print('multi_scale img_size = %g' % dataloader.img_size)
# Update best loss # Update best loss
if rloss['total'] < best_loss: if rloss['total'] < best_loss:
best_loss = rloss['total'] best_loss = rloss['total']
save = True # save training results # Save training results
save = True
if save: if save:
# Save latest checkpoint # Save latest checkpoint
checkpoint = {'epoch': epoch, checkpoint = {'epoch': epoch,
@ -172,7 +175,7 @@ def train(
os.system('cp ' + latest + ' ' + best) os.system('cp ' + latest + ' ' + best)
# Save backup weights every 5 epochs (optional) # Save backup weights every 5 epochs (optional)
if (epoch > 0) & (epoch % 5 == 0): if (epoch > 0) and (epoch % 5 == 0):
os.system('cp ' + latest + ' ' + weights + 'backup{}.pt'.format(epoch)) os.system('cp ' + latest + ' ' + weights + 'backup{}.pt'.format(epoch))
# Calculate mAP # Calculate mAP
@ -188,7 +191,7 @@ if __name__ == '__main__':
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--epochs', type=int, default=270, help='number of epochs') parser.add_argument('--epochs', type=int, default=270, help='number of epochs')
parser.add_argument('--batch-size', type=int, default=16, help='size of each image batch') parser.add_argument('--batch-size', type=int, default=16, help='size of each image batch')
parser.add_argument('--accumulated-batches', type=int, default=1, help='number of batches before optimizer step') parser.add_argument('--accumulate', type=int, default=1, help='accumulate gradient x batches before optimizing')
parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path') parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
parser.add_argument('--data-cfg', type=str, default='cfg/coco.data', help='coco.data file path') parser.add_argument('--data-cfg', type=str, default='cfg/coco.data', help='coco.data file path')
parser.add_argument('--multi-scale', action='store_true', help='random image sizes per batch 320 - 608') parser.add_argument('--multi-scale', action='store_true', help='random image sizes per batch 320 - 608')
@ -206,6 +209,6 @@ if __name__ == '__main__':
resume=opt.resume, resume=opt.resume,
epochs=opt.epochs, epochs=opt.epochs,
batch_size=opt.batch_size, batch_size=opt.batch_size,
accumulated_batches=opt.accumulated_batches, accumulate=opt.accumulate,
multi_scale=opt.multi_scale, multi_scale=opt.multi_scale,
) )

16
utils/datasets.py
View File

@ -90,7 +90,7 @@ class LoadWebcam: # for inference
class LoadImagesAndLabels: # for training class LoadImagesAndLabels: # for training
def __init__(self, path, batch_size=1, img_size=608, multi_scale=False, augment=False): def __init__(self, path, batch_size=1, img_size=608, augment=False):
with open(path, 'r') as file: with open(path, 'r') as file:
self.img_files = file.readlines() self.img_files = file.readlines()
self.img_files = [x.replace('\n', '') for x in self.img_files] self.img_files = [x.replace('\n', '') for x in self.img_files]
@ -102,8 +102,7 @@ class LoadImagesAndLabels: # for training
self.nF = len(self.img_files) # number of image files self.nF = len(self.img_files) # number of image files
self.nB = math.ceil(self.nF / batch_size) # number of batches self.nB = math.ceil(self.nF / batch_size) # number of batches
self.batch_size = batch_size self.batch_size = batch_size
self.height = img_size self.img_size = img_size
self.multi_scale = multi_scale
self.augment = augment self.augment = augment
assert self.nF > 0, 'No images found in %s' % path assert self.nF > 0, 'No images found in %s' % path
@ -121,13 +120,6 @@ class LoadImagesAndLabels: # for training
ia = self.count * self.batch_size ia = self.count * self.batch_size
ib = min((self.count + 1) * self.batch_size, self.nF) ib = min((self.count + 1) * self.batch_size, self.nF)
if self.multi_scale:
# Multi-Scale YOLO Training
height = random.choice(range(10, 20)) * 32 # 320 - 608 pixels
else:
# Fixed-Scale YOLO Training
height = self.height
img_all, labels_all, img_paths, img_shapes = [], [], [], [] img_all, labels_all, img_paths, img_shapes = [], [], [], []
for index, files_index in enumerate(range(ia, ib)): for index, files_index in enumerate(range(ia, ib)):
img_path = self.img_files[self.shuffled_vector[files_index]] img_path = self.img_files[self.shuffled_vector[files_index]]
@ -159,7 +151,7 @@ class LoadImagesAndLabels: # for training
cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img) cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img)
h, w, _ = img.shape h, w, _ = img.shape
img, ratio, padw, padh = letterbox(img, height=height) img, ratio, padw, padh = letterbox(img, height=self.img_size)
# Load labels # Load labels
if os.path.isfile(label_path): if os.path.isfile(label_path):
@ -189,7 +181,7 @@ class LoadImagesAndLabels: # for training
nL = len(labels) nL = len(labels)
if nL > 0: if nL > 0:
# convert xyxy to xywh # convert xyxy to xywh
labels[:, 1:5] = xyxy2xywh(labels[:, 1:5].copy()) / height labels[:, 1:5] = xyxy2xywh(labels[:, 1:5].copy()) / self.img_size
if self.augment: if self.augment:
# random left-right flip # random left-right flip