updates
This commit is contained in:
parent
09d065711a
commit
5e2b802f68
24
detect.py
24
detect.py
|
@ -127,20 +127,18 @@ if __name__ == '__main__':
|
||||||
parser.add_argument('--img-size', type=int, default=416, help='inference size (pixels)')
|
parser.add_argument('--img-size', type=int, default=416, help='inference size (pixels)')
|
||||||
parser.add_argument('--conf-thres', type=float, default=0.5, help='object confidence threshold')
|
parser.add_argument('--conf-thres', type=float, default=0.5, help='object confidence threshold')
|
||||||
parser.add_argument('--nms-thres', type=float, default=0.5, help='iou threshold for non-maximum suppression')
|
parser.add_argument('--nms-thres', type=float, default=0.5, help='iou threshold for non-maximum suppression')
|
||||||
parser.add_argument('--fourcc', type=str, default='mp4v', help='specifies the fourcc code for output video encoding (make sure ffmpeg supports specified fourcc codec)')
|
parser.add_argument('--fourcc', type=str, default='mp4v', help='fourcc output video codec (verify ffmpeg support)')
|
||||||
parser.add_argument('--output', type=str, default='output',help='specifies the output path for images and videos')
|
parser.add_argument('--output', type=str, default='output', help='specifies the output path for images and videos')
|
||||||
opt = parser.parse_args()
|
opt = parser.parse_args()
|
||||||
print(opt)
|
print(opt)
|
||||||
|
|
||||||
with torch.no_grad():
|
with torch.no_grad():
|
||||||
detect(
|
detect(opt.cfg,
|
||||||
opt.cfg,
|
opt.data_cfg,
|
||||||
opt.data_cfg,
|
opt.weights,
|
||||||
opt.weights,
|
images=opt.images,
|
||||||
images=opt.images,
|
img_size=opt.img_size,
|
||||||
img_size=opt.img_size,
|
conf_thres=opt.conf_thres,
|
||||||
conf_thres=opt.conf_thres,
|
nms_thres=opt.nms_thres,
|
||||||
nms_thres=opt.nms_thres,
|
fourcc=opt.fourcc,
|
||||||
fourcc=opt.fourcc,
|
output=opt.output)
|
||||||
output=opt.output
|
|
||||||
)
|
|
||||||
|
|
20
test.py
20
test.py
|
@ -201,14 +201,12 @@ if __name__ == '__main__':
|
||||||
print(opt)
|
print(opt)
|
||||||
|
|
||||||
with torch.no_grad():
|
with torch.no_grad():
|
||||||
mAP = test(
|
mAP = test(opt.cfg,
|
||||||
opt.cfg,
|
opt.data_cfg,
|
||||||
opt.data_cfg,
|
opt.weights,
|
||||||
opt.weights,
|
opt.batch_size,
|
||||||
opt.batch_size,
|
opt.img_size,
|
||||||
opt.img_size,
|
opt.iou_thres,
|
||||||
opt.iou_thres,
|
opt.conf_thres,
|
||||||
opt.conf_thres,
|
opt.nms_thres,
|
||||||
opt.nms_thres,
|
opt.save_json)
|
||||||
opt.save_json
|
|
||||||
)
|
|
||||||
|
|
49
train.py
49
train.py
|
@ -46,12 +46,10 @@ def train(
|
||||||
cfg,
|
cfg,
|
||||||
data_cfg,
|
data_cfg,
|
||||||
img_size=416,
|
img_size=416,
|
||||||
resume=False,
|
|
||||||
epochs=100, # 500200 batches at bs 16, 117263 images = 273 epochs
|
epochs=100, # 500200 batches at bs 16, 117263 images = 273 epochs
|
||||||
batch_size=8,
|
batch_size=8,
|
||||||
accumulate=8, # effective bs = batch_size * accumulate = 8 * 8 = 64
|
accumulate=8, # effective bs = batch_size * accumulate = 8 * 8 = 64
|
||||||
freeze_backbone=False,
|
freeze_backbone=False,
|
||||||
transfer=False # Transfer learning (train only YOLO layers)
|
|
||||||
):
|
):
|
||||||
init_seeds()
|
init_seeds()
|
||||||
weights = 'weights' + os.sep
|
weights = 'weights' + os.sep
|
||||||
|
@ -81,8 +79,8 @@ def train(
|
||||||
start_epoch = 0
|
start_epoch = 0
|
||||||
best_loss = float('inf')
|
best_loss = float('inf')
|
||||||
nf = int(model.module_defs[model.yolo_layers[0] - 1]['filters']) # yolo layer size (i.e. 255)
|
nf = int(model.module_defs[model.yolo_layers[0] - 1]['filters']) # yolo layer size (i.e. 255)
|
||||||
if resume: # Load previously saved model
|
if opt.resume or opt.transfer: # Load previously saved model
|
||||||
if transfer: # Transfer learning
|
if opt.transfer: # Transfer learning
|
||||||
chkpt = torch.load(weights + 'yolov3-spp.pt', map_location=device)
|
chkpt = torch.load(weights + 'yolov3-spp.pt', map_location=device)
|
||||||
model.load_state_dict({k: v for k, v in chkpt['model'].items() if v.numel() > 1 and v.shape[0] != 255},
|
model.load_state_dict({k: v for k, v in chkpt['model'].items() if v.numel() > 1 and v.shape[0] != 255},
|
||||||
strict=False)
|
strict=False)
|
||||||
|
@ -138,7 +136,11 @@ def train(
|
||||||
|
|
||||||
# Initialize distributed training
|
# Initialize distributed training
|
||||||
if torch.cuda.device_count() > 1:
|
if torch.cuda.device_count() > 1:
|
||||||
dist.init_process_group(backend=opt.backend, init_method=opt.dist_url, world_size=opt.world_size, rank=opt.rank)
|
dist.init_process_group(backend='nccl', # 'distributed backend'
|
||||||
|
init_method='tcp://127.0.0.1:9999', # distributed training init method
|
||||||
|
world_size=1, # number of nodes for distributed training
|
||||||
|
rank=0) # distributed training node rank
|
||||||
|
|
||||||
model = torch.nn.parallel.DistributedDataParallel(model)
|
model = torch.nn.parallel.DistributedDataParallel(model)
|
||||||
# sampler = torch.utils.data.distributed.DistributedSampler(dataset)
|
# sampler = torch.utils.data.distributed.DistributedSampler(dataset)
|
||||||
|
|
||||||
|
@ -308,10 +310,6 @@ if __name__ == '__main__':
|
||||||
parser.add_argument('--resume', action='store_true', help='resume training flag')
|
parser.add_argument('--resume', action='store_true', help='resume training flag')
|
||||||
parser.add_argument('--transfer', action='store_true', help='transfer learning flag')
|
parser.add_argument('--transfer', action='store_true', help='transfer learning flag')
|
||||||
parser.add_argument('--num-workers', type=int, default=4, help='number of Pytorch DataLoader workers')
|
parser.add_argument('--num-workers', type=int, default=4, help='number of Pytorch DataLoader workers')
|
||||||
parser.add_argument('--dist-url', default='tcp://127.0.0.1:9999', type=str, help='distributed training init method')
|
|
||||||
parser.add_argument('--rank', default=0, type=int, help='distributed training node rank')
|
|
||||||
parser.add_argument('--world-size', default=1, type=int, help='number of nodes for distributed training')
|
|
||||||
parser.add_argument('--backend', default='nccl', type=str, help='distributed backend')
|
|
||||||
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
|
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
|
||||||
parser.add_argument('--notest', action='store_true', help='only test final epoch')
|
parser.add_argument('--notest', action='store_true', help='only test final epoch')
|
||||||
parser.add_argument('--giou', action='store_true', help='use GIoU loss instead of xy, wh loss')
|
parser.add_argument('--giou', action='store_true', help='use GIoU loss instead of xy, wh loss')
|
||||||
|
@ -325,16 +323,12 @@ if __name__ == '__main__':
|
||||||
opt.nosave = True # only save final checkpoint
|
opt.nosave = True # only save final checkpoint
|
||||||
|
|
||||||
# Train
|
# Train
|
||||||
results = train(
|
results = train(opt.cfg,
|
||||||
opt.cfg,
|
opt.data_cfg,
|
||||||
opt.data_cfg,
|
img_size=opt.img_size,
|
||||||
img_size=opt.img_size,
|
epochs=opt.epochs,
|
||||||
resume=opt.resume or opt.transfer,
|
batch_size=opt.batch_size,
|
||||||
transfer=opt.transfer,
|
accumulate=opt.accumulate)
|
||||||
epochs=opt.epochs,
|
|
||||||
batch_size=opt.batch_size,
|
|
||||||
accumulate=opt.accumulate,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Evolve hyperparameters (optional)
|
# Evolve hyperparameters (optional)
|
||||||
if opt.evolve:
|
if opt.evolve:
|
||||||
|
@ -361,16 +355,12 @@ if __name__ == '__main__':
|
||||||
hyp[k] = np.clip(hyp[k], v[0], v[1])
|
hyp[k] = np.clip(hyp[k], v[0], v[1])
|
||||||
|
|
||||||
# Determine mutation fitness
|
# Determine mutation fitness
|
||||||
results = train(
|
results = train(opt.cfg,
|
||||||
opt.cfg,
|
opt.data_cfg,
|
||||||
opt.data_cfg,
|
img_size=opt.img_size,
|
||||||
img_size=opt.img_size,
|
epochs=opt.epochs,
|
||||||
resume=opt.resume or opt.transfer,
|
batch_size=opt.batch_size,
|
||||||
transfer=opt.transfer,
|
accumulate=opt.accumulate)
|
||||||
epochs=opt.epochs,
|
|
||||||
batch_size=opt.batch_size,
|
|
||||||
accumulate=opt.accumulate,
|
|
||||||
)
|
|
||||||
mutation_fitness = results[2]
|
mutation_fitness = results[2]
|
||||||
|
|
||||||
# Write mutation results
|
# Write mutation results
|
||||||
|
@ -378,7 +368,6 @@ if __name__ == '__main__':
|
||||||
|
|
||||||
# Update hyperparameters if fitness improved
|
# Update hyperparameters if fitness improved
|
||||||
if mutation_fitness > best_fitness:
|
if mutation_fitness > best_fitness:
|
||||||
# Fitness improved!
|
|
||||||
print('Fitness improved!')
|
print('Fitness improved!')
|
||||||
best_fitness = mutation_fitness
|
best_fitness = mutation_fitness
|
||||||
else:
|
else:
|
||||||
|
|
Loading…
Reference in New Issue