def train(train_loop_func, logger, args):
if args.amp:
amp_handle = amp.init(enabled=args.fp16)
# Check that GPUs are actually available
use_cuda = not args.no_cuda
# Setup multi-GPU if necessary
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
if args.distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.N_gpu = torch.distributed.get_world_size()
else:
args.N_gpu = 1
if args.seed is None:
args.seed = np.random.randint(1e4)
if args.distributed:
args.seed = (args.seed + torch.distributed.get_rank()) % 2**32
print("Using seed = {}".format(args.seed))
torch.manual_seed(args.seed)
np.random.seed(seed=args.seed)
# Setup data, defaults
dboxes = dboxes300_coco()
encoder = Encoder(dboxes)
cocoGt = get_coco_ground_truth(args)
train_loader = get_train_loader(args, args.seed - 2**31)
val_dataset = get_val_dataset(args)
val_dataloader = get_val_dataloader(val_dataset, args)
ssd300 = SSD300(backbone=args.backbone)
args.learning_rate = args.learning_rate * args.N_gpu * (args.batch_size /
32)
start_epoch = 0
iteration = 0
loss_func = Loss(dboxes)
if use_cuda:
ssd300.cuda()
loss_func.cuda()
if args.fp16 and not args.amp:
ssd300 = network_to_half(ssd300)
if args.distributed:
ssd300 = DDP(ssd300)
optimizer = torch.optim.SGD(tencent_trick(ssd300),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
scheduler = MultiStepLR(optimizer=optimizer,
milestones=args.multistep,
gamma=0.1)
if args.fp16:
if args.amp:
optimizer = amp_handle.wrap_optimizer(optimizer)
else:
optimizer = FP16_Optimizer(optimizer, static_loss_scale=128.)
if args.checkpoint is not None:
if os.path.isfile(args.checkpoint):
load_checkpoint(ssd300, args.checkpoint)
checkpoint = torch.load(args.checkpoint,
map_location=lambda storage, loc: storage.
cuda(torch.cuda.current_device()))
start_epoch = checkpoint['epoch']
iteration = checkpoint['iteration']
scheduler.load_state_dict(checkpoint['scheduler'])
ssd300.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
else:
print('Provided checkpoint is not path to a file')
return
inv_map = {v: k for k, v in val_dataset.label_map.items()}
total_time = 0
if args.mode == 'evaluation':
acc = evaluate(ssd300, val_dataloader, cocoGt, encoder, inv_map, args)
if args.local_rank == 0:
print('Model precision {} mAP'.format(acc))
return
mean, std = generate_mean_std(args)
for epoch in range(start_epoch, args.epochs):
start_epoch_time = time.time()
scheduler.step()
iteration = train_loop_func(ssd300, loss_func, epoch, optimizer,
train_loader, val_dataloader, encoder,
iteration, logger, args, mean, std)
end_epoch_time = time.time() - start_epoch_time
total_time += end_epoch_time
if args.local_rank == 0:
logger.update_epoch_time(epoch, end_epoch_time)
if epoch in args.evaluation:
acc = evaluate(ssd300, val_dataloader, cocoGt, encoder, inv_map,
args)
if args.local_rank == 0:
logger.update_epoch(epoch, acc)
if args.save and args.local_rank == 0:
print("saving model...")
obj = {
'epoch': epoch + 1,
'iteration': iteration,
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict(),
'label_map': val_dataset.label_info
}
if args.distributed:
obj['model'] = ssd300.module.state_dict()
else:
obj['model'] = ssd300.state_dict()
torch.save(obj, './models/epoch_{}.pt'.format(epoch))
train_loader.reset()
print('total training time: {}'.format(total_time))
def train(train_loop_func, logger, args):
# Check that GPUs are actually available
use_cuda = not args.no_cuda
# Setup multi-GPU if necessary
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
if args.distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.N_gpu = torch.distributed.get_world_size()
else:
args.N_gpu = 1
if args.seed is None:
args.seed = np.random.randint(1e4)
if args.distributed:
args.seed = (args.seed + torch.distributed.get_rank()) % 2**32
print("Using seed = {}".format(args.seed))
torch.manual_seed(args.seed)
np.random.seed(seed=args.seed)
# Setup data, defaults
dboxes = dboxes300_coco()
encoder = Encoder(dboxes)
cocoGt = get_coco_ground_truth(args)
train_loader = get_train_loader(args, args.seed - 2**31)
val_dataset = get_val_dataset(args)
val_dataloader = get_val_dataloader(val_dataset, args)
ssd300 = SSD300(backbone=ResNet(args.backbone, args.backbone_path))
# args.learning_rate = args.learning_rate * args.N_gpu * (args.batch_size / 32)
print(f"Actual starting LR: {args.learning_rate}")
start_epoch = 0
iteration = 0
loss_func = Loss(dboxes)
if use_cuda:
ssd300.cuda()
loss_func.cuda()
# optimizer = torch.optim.SGD(tencent_trick(ssd300), lr=args.learning_rate,
# momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True)
optimizer = torch.optim.AdamW(tencent_trick(ssd300),
lr=args.learning_rate,
betas=(0.8, 0.999),
eps=1e-08,
weight_decay=0.01,
amsgrad=True)
# scheduler = MultiStepLR(optimizer=optimizer, milestones=args.multistep, gamma=0.1)
# scheduler = CosineAnnealingWarmRestarts(optimizer=optimizer, T_0=20, T_mult=1, eta_min=1e-6)
scheduler = CosineAnnealingLR(optimizer=optimizer,
T_max=args.epochs,
eta_min=1e-6)
# scheduler = OneCycleLR(optimizer, max_lr=0.003, epochs=41, steps_per_epoch=173)
# scheduler = CyclicLR(optimizer, base_lr=args.learning_rate, max_lr=2*args.learning_rate,
# step_size_up=173*3, step_size_down=173*10)
if args.amp:
ssd300, optimizer = amp.initialize(ssd300, optimizer, opt_level='O2')
if args.distributed:
ssd300 = DDP(ssd300)
if args.checkpoint is not None:
if os.path.isfile(args.checkpoint):
load_checkpoint(ssd300.module if args.distributed else ssd300,
args.checkpoint)
checkpoint = torch.load(args.checkpoint,
map_location=lambda storage, loc: storage.
cuda(torch.cuda.current_device()))
start_epoch = checkpoint['epoch']
iteration = checkpoint['iteration']
scheduler.load_state_dict(checkpoint['scheduler'])
optimizer.load_state_dict(checkpoint['optimizer'])
else:
print('Provided checkpoint is not path to a file')
return
inv_map = {v: k for k, v in val_dataset.label_map.items()}
total_time = 0
if args.mode == 'evaluation':
acc = evaluate(ssd300, val_dataloader, cocoGt, encoder, inv_map, args)
if args.local_rank == 0:
print('Model precision {} mAP'.format(acc))
return
mean, std = generate_mean_std(args)
for epoch in range(start_epoch, args.epochs):
start_epoch_time = time.time()
# scheduler.step()
iteration = train_loop_func(ssd300, loss_func, epoch, optimizer,
scheduler, train_loader, val_dataloader,
encoder, iteration, logger, args, mean,
std)
end_epoch_time = time.time() - start_epoch_time
total_time += end_epoch_time
# https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate
scheduler.step()
if args.local_rank == 0:
logger.update_epoch_time(epoch, end_epoch_time)
if epoch in args.evaluation:
acc = evaluate(ssd300, val_dataloader, cocoGt, encoder, inv_map,
args)
if args.local_rank == 0:
logger.update_epoch(epoch, acc)
if args.save and args.local_rank == 0:
print("saving model...")
obj = {
'epoch': epoch + 1,
'iteration': iteration,
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict(),
'label_map': val_dataset.label_info
}
if args.distributed:
obj['model'] = ssd300.module.state_dict()
else:
obj['model'] = ssd300.state_dict()
torch.save(obj, './models/epoch_{}.pt'.format(epoch))
train_loader.reset()
print('total training time: {}'.format(total_time))
def train(train_loop_func, logger, args):
# Check that GPUs are actually available
use_cuda = not args.no_cuda
train_samples = 118287
# Setup multi-GPU if necessary
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
if args.distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='smddp', init_method='env://')
args.N_gpu = torch.distributed.get_world_size()
else:
args.N_gpu = 1
if args.seed is None:
args.seed = np.random.randint(1e4)
if args.distributed:
args.seed = (args.seed + torch.distributed.get_rank()) % 2**32
print("Using seed = {}".format(args.seed))
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
np.random.seed(seed=args.seed)
# Setup data, defaults
dboxes = dboxes300_coco()
encoder = Encoder(dboxes)
cocoGt = get_coco_ground_truth(args)
train_loader = get_train_loader(args, args.seed - 2**31)
val_dataset = get_val_dataset(args)
val_dataloader = get_val_dataloader(val_dataset, args)
ssd300 = SSD300(backbone=ResNet(args.backbone, args.backbone_path))
args.learning_rate = args.learning_rate * args.N_gpu * (args.batch_size / 32)
start_epoch = 0
iteration = 0
loss_func = Loss(dboxes)
if use_cuda:
ssd300.cuda()
loss_func.cuda()
optimizer = torch.optim.SGD(tencent_trick(ssd300), lr=args.learning_rate,
momentum=args.momentum, weight_decay=args.weight_decay)
scheduler = MultiStepLR(optimizer=optimizer, milestones=args.multistep, gamma=0.1)
if args.amp:
ssd300, optimizer = amp.initialize(ssd300, optimizer, opt_level='O2')
if args.distributed:
ssd300 = DDP(ssd300)
if args.checkpoint is not None:
if os.path.isfile(args.checkpoint):
load_checkpoint(ssd300.module if args.distributed else ssd300, args.checkpoint)
checkpoint = torch.load(args.checkpoint,
map_location=lambda storage, loc: storage.cuda(torch.cuda.current_device()))
start_epoch = checkpoint['epoch']
iteration = checkpoint['iteration']
scheduler.load_state_dict(checkpoint['scheduler'])
optimizer.load_state_dict(checkpoint['optimizer'])
else:
print('Provided checkpoint is not path to a file')
return
inv_map = {v: k for k, v in val_dataset.label_map.items()}
total_time = 0
if args.mode == 'evaluation':
acc = evaluate(ssd300, val_dataloader, cocoGt, encoder, inv_map, args)
if args.local_rank == 0:
print('Model precision {} mAP'.format(acc))
return
mean, std = generate_mean_std(args)
for epoch in range(start_epoch, args.epochs):
start_epoch_time = time.time()
scheduler.step()
iteration = train_loop_func(ssd300, loss_func, epoch, optimizer, train_loader, val_dataloader, encoder, iteration,
logger, args, mean, std)
end_epoch_time = time.time() - start_epoch_time
total_time += end_epoch_time
if torch.distributed.get_rank() == 0:
throughput = train_samples / end_epoch_time
logger.update_epoch_time(epoch, end_epoch_time)
logger.update_throughput_speed(epoch, throughput)
if epoch in args.evaluation:
acc = evaluate(ssd300, val_dataloader, cocoGt, encoder, inv_map, args)
if args.save and args.local_rank == 0:
print("saving model...")
obj = {'epoch': epoch + 1,
'iteration': iteration,
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict(),
'label_map': val_dataset.label_info}
if args.distributed:
obj['model'] = ssd300.module.state_dict()
else:
obj['model'] = ssd300.state_dict()
save_path = os.path.join(args.save, f'epoch_{epoch}.pt')
torch.save(obj, save_path)
logger.log('model path', save_path)
train_loader.reset()
if torch.distributed.get_rank() == 0:
DLLogger.log((), { 'Total training time': '%.2f' % total_time + ' secs' })
logger.log_summary()
def train(train_loop_func, args, logger):
# Setup multi-GPU if necessary
# args.distributed = False
# if 'WORLD_SIZE' in os.environ:
# args.distributed = int(os.environ['WORLD_SIZE']) > 1
# args.distributed = True
# if args.distributed:
# torch.cuda.set_device(args.local_rank)
# torch.distributed.init_process_group(backend='nccl')
# args.N_gpu = torch.distributed.get_world_size()
# else:
# args.N_gpu = 1
if args.seed is None:
args.seed = np.random.randint(10000)
# if args.distributed:
# args.seed = (args.seed + torch.distributed.get_rank()) % 2 ** 32
print("Using seed = {}".format(args.seed))
torch.manual_seed(args.seed)
np.random.seed(seed=args.seed)
random.seed(args.seed)
torch.backends.cudnn.deterministic = True
model = EfficientNet.from_pretrained('efficientnet-b4', num_classes=4)
if args.local_rank is not None:
torch.distributed.init_process_group(backend="nccl")
torch.cuda.set_device(args.local_rank)
model = model.cuda()
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
{'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
'weight_decay': 0.002},
{'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
'weight_decay': 0.0}
]
optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=args.learning_rate)
if args.amp:
model, optimizer = amp.initialize(model, optimizer, opt_level='O2')
if args.local_rank is not None:
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
output_device=args.local_rank)
# Setup data, defaults
train, test = construct_dataset(args.data)
random.shuffle(train)
train = train[:len(train)]
split_position = int(len(train) * 0.96)
train_dataset = ALASKA2Dataset(train[:split_position], root_dir=args.data, augmented=True)
val_dataset = ALASKA2Dataset(train[split_position:], root_dir=args.data, augmented=False)
test_dataset = ALASKA2Dataset(test, root_dir=args.data, augmented=False)
if args.local_rank is not None:
train_sampler = DistributedSampler(dataset=train_dataset, shuffle=True)
train_sampler.set_epoch(0)
else:
train_sampler = RandomSampler(train_dataset)
train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, drop_last=False,
num_workers=4, shuffle=False, sampler=train_sampler)
val_dataloader = DataLoader(val_dataset, batch_size=args.eval_batch_size, drop_last=False,
num_workers=4, shuffle=False)
test_dataloader = DataLoader(test_dataset, batch_size=args.eval_batch_size, drop_last=False,
num_workers=4, shuffle=False)
mean, std = generate_mean_std(amp=args.amp)
# args.learning_rate = args.learning_rate * args.N_gpu * (args.batch_size / 32)
# scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.5, patience=2,
verbose=False, threshold_mode='abs')
start_epoch = 1
if args.checkpoint is not None:
if os.path.isfile(args.checkpoint):
model, optimizer, scheduler, start_epoch = load_checkpoint(
args.checkpoint, model, optimizer, scheduler)
start_epoch += 1 # this is because the epoch saved is the previous epoch
else:
print('Provided checkpoint is not path to a file')
return
# loss_function = nn.CrossEntropyLoss()
loss_function = LabelSmoothing()
if args.mode == 'evaluation':
acc = evaluate(model, val_dataloader, args, mean, std, loss_function)
print('Model precision {} mAP'.format(acc))
return
elif args.mode == 'testing':
test_(model, test_dataloader, args, mean, std)
return
for epoch in range(start_epoch, args.epochs + 1):
print("-----------------------")
print("Local Rank: {}, Epoch: {}, Training ...".format(args.local_rank, epoch))
print("Epoch {} of {}".format(epoch, args.epochs))
print("Total number of parameters trained this epoch: ",
sum(p.numel() for pg in optimizer.param_groups for p in pg['params'] if
p.requires_grad))
avg_loss = train_loop_func(model, loss_function, optimizer, train_dataloader, None, args,
mean, std)
# logger.update_epoch_time(epoch, end_epoch_time)
print("saving model...")
obj = {'epoch': epoch,
'model': model.module.state_dict(), # model.state_dict() for non DataParallel model
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict()}
if args.local_rank in [0, None]:
torch.save(obj, f'./saved/{args.backbone}_epoch_{epoch}.pt')
print("Incepe evaluarea")
val_loss = evaluate(model, val_dataloader, args, mean, std, loss_function)
test_(model, test_dataloader, args, mean, std, epoch)
scheduler.step(val_loss)
