def test(model, device):
global cfg
if args.testset:
print('test on test-dev 2017')
evaluator = COCOAPIEvaluator(data_dir=data_dir,
img_size=cfg['min_dim'],
device=device,
testset=True,
transform=BaseTransform(
cfg['min_dim'],
mean=(0.406, 0.456, 0.485),
std=(0.225, 0.224, 0.229)))
else:
# eval
evaluator = COCOAPIEvaluator(data_dir=data_dir,
img_size=cfg['min_dim'],
device=device,
testset=False,
transform=BaseTransform(
cfg['min_dim'],
mean=(0.406, 0.456, 0.485),
std=(0.225, 0.224, 0.229)))
# COCO evaluation
ap50_95, ap50 = evaluator.evaluate(model)
print('ap50 : ', ap50)
print('ap50_95 : ', ap50_95)
def main():
"""
YOLOv3 trainer. See README for details.
"""
args = parse_args()
print("Setting Arguments.. : ", args)
cuda = torch.cuda.is_available() and args.use_cuda
os.makedirs(args.checkpoint_dir, exist_ok=True)
# Parse config settings
with open(args.cfg, 'r') as f:
cfg = yaml.load(f)
print("successfully loaded config file: ", cfg)
momentum = cfg['TRAIN']['MOMENTUM']
decay = cfg['TRAIN']['DECAY']
burn_in = cfg['TRAIN']['BURN_IN']
iter_size = cfg['TRAIN']['MAXITER']
steps = eval(cfg['TRAIN']['STEPS'])
batch_size = cfg['TRAIN']['BATCHSIZE']
subdivision = cfg['TRAIN']['SUBDIVISION']
ignore_thre = cfg['TRAIN']['IGNORETHRE']
random_resize = cfg['AUGMENTATION']['RANDRESIZE']
base_lr = cfg['TRAIN']['LR'] / batch_size / subdivision
print('effective_batch_size = batch_size * iter_size = %d * %d' %
(batch_size, subdivision))
# Initiate model
model = YOLOv3(cfg['MODEL'], ignore_thre=ignore_thre)
if args.weights_path:
print("loading darknet weights....", args.weights_path)
parse_yolo_weights(model, args.weights_path)
elif args.checkpoint:
print("loading pytorch ckpt...", args.checkpoint)
state = torch.load(args.checkpoint)
if 'model_state_dict' in state.keys():
model.load_state_dict(state['model_state_dict'])
else:
model.load_state_dict(state)
if cuda:
print("using cuda")
model = model.cuda()
model.train()
imgsize = cfg['TRAIN']['IMGSIZE']
evaluator = COCOAPIEvaluator(model_type=cfg['MODEL']['TYPE'],
data_dir='COCO/',
img_size=cfg['TEST']['IMGSIZE'],
confthre=cfg['TEST']['CONFTHRE'],
nmsthre=cfg['TEST']['NMSTHRE'])
ap50_95, ap50 = evaluator.evaluate(model)
def test(model, device):
global cfg
evaluator = COCOAPIEvaluator(data_dir=data_dir,
img_size=cfg['min_dim'],
device=device,
transform=BaseTransform(cfg['min_dim'],
mean=(0.406, 0.456,
0.485),
std=(0.225, 0.224,
0.229)))
# COCO evaluation
ap50_95, ap50 = evaluator.evaluate(model)
print('ap50 : ', ap50)
print('ap50_95 : ', ap50_95)
def test(model, device, input_size):
if args.testset:
print('test on test-dev 2017')
evaluator = COCOAPIEvaluator(data_dir=data_dir,
img_size=input_size,
device=device,
testset=True,
transform=BaseTransform(input_size))
else:
# eval
evaluator = COCOAPIEvaluator(data_dir=data_dir,
img_size=input_size,
device=device,
testset=False,
transform=BaseTransform(input_size))
# COCO evaluation
ap50_95, ap50 = evaluator.evaluate(model)
print('ap50 : ', ap50)
print('ap50_95 : ', ap50_95)
def coco_test(model, device, input_size, test=False):
if test:
# test-dev
print('test on test-dev 2017')
evaluator = COCOAPIEvaluator(data_dir=coco_root,
img_size=input_size,
device=device,
testset=True,
transform=BaseTransform(input_size))
else:
# eval
evaluator = COCOAPIEvaluator(data_dir=coco_root,
img_size=input_size,
device=device,
testset=False,
transform=BaseTransform(input_size))
# COCO evaluation
evaluator.evaluate(model)
def main():
"""
YOLOv3 trainer. See README for details.
"""
args = parse_args()
print("Setting Arguments.. : ", args)
cuda = torch.cuda.is_available() and args.use_cuda
os.makedirs(args.log_dir, exist_ok=True)
os.makedirs(args.save_dir, exist_ok=True)
if args.distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend="nccl",
init_method="env://")
save_prefix = 'yolov3'
# Parse config settings
with open(args.cfg, 'r') as f:
cfg = yaml.safe_load(f)
print("successfully loaded config file: ", cfg)
lr = cfg['TRAIN']['LR']
epochs = cfg['TRAIN']['MAXEPOCH']
cos = cfg['TRAIN']['COS']
sybn = cfg['TRAIN']['SYBN']
mixup = cfg['TRAIN']['MIX']
no_mixup_epochs = cfg['TRAIN']['NO_MIXUP_EPOCHS']
label_smooth = cfg['TRAIN']['LABAL_SMOOTH']
momentum = cfg['TRAIN']['MOMENTUM']
burn_in = cfg['TRAIN']['BURN_IN']
batch_size = cfg['TRAIN']['BATCHSIZE']
decay = cfg['TRAIN']['DECAY']
ignore_thre = cfg['TRAIN']['IGNORETHRE']
random_resize = cfg['TRAIN']['RANDRESIZE']
input_size = (cfg['TRAIN']['IMGSIZE'], cfg['TRAIN']['IMGSIZE'])
test_size = (args.test_size, args.test_size)
step = (180, 240) # for no cos lr shedule training
# Learning rate setup
base_lr = lr
if args.dataset == 'COCO':
dataset = COCODataset(data_dir='data/COCO/',
img_size=input_size,
preproc=TrainTransform(rgb_means=(0.485, 0.456,
0.406),
std=(0.229, 0.224, 0.225),
max_labels=50),
debug=args.debug)
num_class = 80
elif args.dataset == 'VOC':
train_sets = [('2007', 'trainval'), ('2012', 'trainval')]
dataset = VOCDetection(root='data/VOC',
image_sets=train_sets,
input_dim=input_size,
preproc=TrainTransform(rgb_means=(0.485, 0.456,
0.406),
std=(0.229, 0.224,
0.225),
max_labels=30))
num_class = 20
else:
print('Only COCO and VOC datasets are supported!')
return
save_prefix += ('_' + args.dataset)
if label_smooth:
save_prefix += '_label_smooth'
# Initiate model
if args.asff:
save_prefix += '_asff'
from models.yolov3_asff import YOLOv3
print('Training YOLOv3 with ASFF!')
model = YOLOv3(num_classes=num_class,
ignore_thre=ignore_thre,
label_smooth=label_smooth,
rfb=args.rfb,
vis=args.vis)
else:
save_prefix += '_baseline'
from models.yolov3_baseline import YOLOv3
print('Training YOLOv3 strong baseline!')
if args.vis:
print('Visualization is not supported for YOLOv3 baseline model')
args.vis = False
model = YOLOv3(num_classes=num_class,
ignore_thre=ignore_thre,
label_smooth=label_smooth,
rfb=args.rfb)
save_to_disk = (not args.distributed) or distributed_util.get_rank() == 0
def init_yolo(m):
for key in m.state_dict():
if key.split('.')[-1] == 'weight':
if 'conv' in key:
init.kaiming_normal_(m.state_dict()[key],
a=0.1,
mode='fan_in')
if 'linear' in key:
init.kaiming_normal_(m.state_dict()[key],
a=0.0,
mode='fan_in')
if 'bn' in key:
m.state_dict()[key][...] = 1
elif key.split('.')[-1] == 'bias':
m.state_dict()[key][...] = 0
model.apply(init_yolo)
if sybn:
model = apex.parallel.convert_syncbn_model(model)
if args.checkpoint:
print("loading pytorch ckpt...", args.checkpoint)
cpu_device = torch.device("cpu")
ckpt = torch.load(args.checkpoint, map_location=cpu_device)
model.load_state_dict(ckpt, strict=False)
#model.load_state_dict(ckpt)
if cuda:
print("using cuda")
torch.backends.cudnn.benchmark = True
device = torch.device("cuda")
model = model.to(device)
if args.half:
model = model.half()
if args.ngpu > 1:
if args.distributed:
model = apex.parallel.DistributedDataParallel(model,
delay_allreduce=True)
#model = apex.parallel.DistributedDataParallel(model)
else:
model = nn.DataParallel(model)
if args.tfboard and save_to_disk:
print("using tfboard")
from torch.utils.tensorboard import SummaryWriter
tblogger = SummaryWriter(args.log_dir)
model.train()
if mixup:
from dataset.mixupdetection import MixupDetection
dataset = MixupDetection(
dataset,
preproc=TrainTransform(rgb_means=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
max_labels=50),
)
dataset.set_mixup(np.random.beta, 1.5, 1.5)
save_prefix += '_mixup'
if args.distributed:
sampler = torch.utils.data.DistributedSampler(dataset)
else:
sampler = torch.utils.data.RandomSampler(dataset)
batch_sampler = YoloBatchSampler(sampler=sampler,
batch_size=batch_size,
drop_last=False,
input_dimension=input_size)
dataloader = DataLoader(dataset,
batch_sampler=batch_sampler,
num_workers=args.n_cpu,
pin_memory=True)
dataiterator = iter(dataloader)
if args.dataset == 'COCO':
evaluator = COCOAPIEvaluator(data_dir='data/COCO/',
img_size=test_size,
confthre=cfg['TEST']['CONFTHRE'],
nmsthre=cfg['TEST']['NMSTHRE'],
testset=args.testset,
vis=args.vis)
elif args.dataset == 'VOC':
'''
# COCO style evaluation, you have to convert xml annotation files into a json file.
evaluator = COCOAPIEvaluator(
data_dir='data/VOC/',
img_size=test_size,
confthre=cfg['TEST']['CONFTHRE'],
nmsthre=cfg['TEST']['NMSTHRE'],
testset=args.testset,
voc = True)
'''
evaluator = VOCEvaluator(data_dir='data/VOC/',
img_size=test_size,
confthre=cfg['TEST']['CONFTHRE'],
nmsthre=cfg['TEST']['NMSTHRE'],
vis=args.vis)
dtype = torch.float16 if args.half else torch.float32
# optimizer setup
# set weight decay only on conv.weight
if args.no_wd:
params_dict = dict(model.named_parameters())
params = []
for key, value in params_dict.items():
if 'conv.weight' in key:
params += [{'params': value, 'weight_decay': decay}]
else:
params += [{'params': value, 'weight_decay': 0.0}]
save_prefix += '_no_wd'
else:
params = model.parameters()
optimizer = optim.SGD(params,
lr=base_lr,
momentum=momentum,
dampening=0,
weight_decay=decay)
if args.half:
optimizer = FP16_Optimizer(optimizer, verbose=False)
if cos:
save_prefix += '_cos'
tmp_lr = base_lr
def set_lr(tmp_lr):
for param_group in optimizer.param_groups:
param_group['lr'] = tmp_lr
# start training loop
start = time.time()
epoch = args.start_epoch
epoch_size = len(dataset) // (batch_size * args.ngpu)
while epoch < epochs + 1:
if args.distributed:
batch_sampler.sampler.set_epoch(epoch)
if epoch > epochs - no_mixup_epochs + 1:
args.eval_interval = 1
if mixup:
print('Disable mix up now!')
mixup = False
dataset.set_mixup(None)
if args.distributed:
sampler = torch.utils.data.DistributedSampler(dataset)
else:
sampler = torch.utils.data.RandomSampler(dataset)
batch_sampler = YoloBatchSampler(sampler=sampler,
batch_size=batch_size,
drop_last=False,
input_dimension=input_size)
dataloader = DataLoader(dataset,
batch_sampler=batch_sampler,
num_workers=args.n_cpu,
pin_memory=True)
#### DropBlock Shedule #####
Drop_layer = [16, 24, 33]
if args.asff:
Drop_layer = [16, 22, 29]
if (epoch == 5 or (epoch == args.start_epoch
and args.start_epoch > 5)) and (args.dropblock):
block_size = [1, 3, 5]
keep_p = [0.9, 0.9, 0.9]
for i in range(len(Drop_layer)):
model.module.module_list[Drop_layer[i]].reset(
block_size[i], keep_p[i])
if (epoch == 80 or (epoch == args.start_epoch
and args.start_epoch > 80)) and (args.dropblock):
block_size = [3, 5, 7]
keep_p = [0.9, 0.9, 0.9]
for i in range(len(Drop_layer)):
model.module.module_list[Drop_layer[i]].reset(
block_size[i], keep_p[i])
if (epoch == 150 or (epoch == args.start_epoch
and args.start_epoch > 150)) and (args.dropblock):
block_size = [7, 7, 7]
keep_p = [0.9, 0.9, 0.9]
for i in range(len(Drop_layer)):
model.module.module_list[Drop_layer[i]].reset(
block_size[i], keep_p[i])
for iter_i, (imgs, targets, img_info, idx) in enumerate(dataloader):
#evaluation
if ((epoch % args.eval_interval == 0) and epoch > args.start_epoch
and iter_i == 0) or args.test:
if not args.test and save_to_disk:
torch.save(
model.module.state_dict(),
os.path.join(args.save_dir,
save_prefix + '_' + repr(epoch) + '.pth'))
if args.distributed:
distributed_util.synchronize()
ap50_95, ap50 = evaluator.evaluate(model, args.half)
if args.distributed:
distributed_util.synchronize()
if args.test:
sys.exit(0)
model.train()
if args.tfboard and save_to_disk:
tblogger.add_scalar('val/COCOAP50', ap50, epoch)
tblogger.add_scalar('val/COCOAP50_95', ap50_95, epoch)
# learning rate scheduling (cos or step)
if epoch < burn_in:
tmp_lr = base_lr * pow((iter_i + epoch * epoch_size) * 1. /
(burn_in * epoch_size), 3)
set_lr(tmp_lr)
elif cos:
if epoch <= epochs - no_mixup_epochs and epoch > 20:
min_lr = 0.00001
tmp_lr = min_lr + 0.5*(base_lr-min_lr)*(1+math.cos(math.pi*(epoch-20)*1./\
(epochs-no_mixup_epochs-20)))
elif epoch > epochs - no_mixup_epochs:
tmp_lr = 0.00001
set_lr(tmp_lr)
elif epoch == burn_in:
tmp_lr = base_lr
set_lr(tmp_lr)
elif epoch in steps and iter_i == 0:
tmp_lr = tmp_lr * 0.1
set_lr(tmp_lr)
optimizer.zero_grad()
imgs = Variable(imgs.to(device).to(dtype))
targets = Variable(targets.to(device).to(dtype),
requires_grad=False)
loss_dict = model(imgs, targets, epoch)
loss_dict_reduced = reduce_loss_dict(loss_dict)
loss = sum(loss for loss in loss_dict['losses'])
if args.half:
optimizer.backward(loss)
else:
loss.backward()
#torch.nn.utils.clip_grad_norm_(model.parameters(), 10)
optimizer.step()
if iter_i % 10 == 0 and save_to_disk:
# logging
end = time.time()
print(
'[Epoch %d/%d][Iter %d/%d][lr %.6f]'
'[Loss: anchor %.2f, iou %.2f, l1 %.2f, conf %.2f, cls %.2f, imgsize %d, time: %.2f]'
%
(epoch, epochs, iter_i, epoch_size, tmp_lr,
sum(anchor_loss for anchor_loss in
loss_dict_reduced['anchor_losses']).item(),
sum(iou_loss for iou_loss in
loss_dict_reduced['iou_losses']).item(),
sum(l1_loss
for l1_loss in loss_dict_reduced['l1_losses']).item(),
sum(conf_loss for conf_loss in
loss_dict_reduced['conf_losses']).item(),
sum(cls_loss
for cls_loss in loss_dict_reduced['cls_losses']).item(
), input_size[0], end - start),
flush=True)
start = time.time()
if args.tfboard and save_to_disk:
tblogger.add_scalar(
'train/total_loss',
sum(loss
for loss in loss_dict_reduced['losses']).item(),
epoch * epoch_size + iter_i)
# random resizing
if random_resize and iter_i % 10 == 0 and iter_i > 0:
tensor = torch.LongTensor(1).to(device)
if args.distributed:
distributed_util.synchronize()
if save_to_disk:
if epoch > epochs - 10:
size = 416 if args.dataset == 'VOC' else 608
else:
size = random.randint(*(10, 19))
size = int(32 * size)
tensor.fill_(size)
if args.distributed:
distributed_util.synchronize()
dist.broadcast(tensor, 0)
input_size = dataloader.change_input_dim(
multiple=tensor.item(), random_range=None)
if args.distributed:
distributed_util.synchronize()
epoch += 1
if not args.test and save_to_disk:
torch.save(
model.module.state_dict(),
os.path.join(args.save_dir,
"yolov3_" + args.dataset + '_Final.pth'))
if args.distributed:
distributed_util.synchronize()
ap50_95, ap50 = evaluator.evaluate(model, args.half)
if args.tfboard and save_to_disk:
tblogger.close()
def train():
args = parse_args()
data_dir = coco_root
path_to_save = os.path.join(args.save_folder, args.version)
os.makedirs(path_to_save, exist_ok=True)
hr = False
if args.high_resolution:
print('use hi-res backbone')
hr = True
cfg = coco_af
if args.cuda:
print('use cuda')
cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
if args.mosaic:
print("use Mosaic Augmentation ...")
# multi scale
if args.multi_scale:
print('Let us use the multi-scale trick.')
input_size = [640, 640]
else:
input_size = [416, 416]
print("Setting Arguments.. : ", args)
print("----------------------------------------------------------")
print('Loading the MSCOCO dataset...')
# dataset
dataset = COCODataset(data_dir=data_dir,
img_size=input_size[0],
transform=SSDAugmentation(input_size),
debug=args.debug,
mosaic=args.mosaic)
# build model
if args.version == 'yolo':
from models.yolo import myYOLO
yolo_net = myYOLO(device,
input_size=input_size,
num_classes=args.num_classes,
trainable=True,
hr=hr)
print('Let us train yolo on the COCO dataset ......')
else:
print('We only support YOLO !!!')
exit()
print("----------------------------------------------------------")
print('The dataset size:', len(dataset))
print("----------------------------------------------------------")
# use tfboard
if args.tfboard:
print('use tensorboard')
from torch.utils.tensorboard import SummaryWriter
c_time = time.strftime('%Y-%m-%d %H:%M:%S',
time.localtime(time.time()))
log_path = os.path.join('log/coco/', args.version, c_time)
os.makedirs(log_path, exist_ok=True)
writer = SummaryWriter(log_path)
model = yolo_net
model.to(device).train()
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
collate_fn=detection_collate,
num_workers=args.num_workers)
evaluator = COCOAPIEvaluator(data_dir=data_dir,
img_size=cfg['min_dim'],
device=device,
transform=BaseTransform(cfg['min_dim']))
# optimizer setup
base_lr = args.lr
tmp_lr = base_lr
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
max_epoch = cfg['max_epoch']
epoch_size = len(dataset) // args.batch_size
# start training loop
t0 = time.time()
for epoch in range(max_epoch):
# use cos lr
if args.cos and epoch > 20 and epoch <= max_epoch - 20:
# use cos lr
tmp_lr = 0.00001 + 0.5 * (base_lr - 0.00001) * (
1 + math.cos(math.pi * (epoch - 20) * 1. / (max_epoch - 20)))
set_lr(optimizer, tmp_lr)
elif args.cos and epoch > max_epoch - 20:
tmp_lr = 0.00001
set_lr(optimizer, tmp_lr)
# use step lr
else:
if epoch in cfg['lr_epoch']:
tmp_lr = tmp_lr * 0.1
set_lr(optimizer, tmp_lr)
for iter_i, (images, targets) in enumerate(dataloader):
# WarmUp strategy for learning rate
if not args.no_warm_up:
if epoch < args.wp_epoch:
tmp_lr = base_lr * pow((iter_i + epoch * epoch_size) * 1. /
(args.wp_epoch * epoch_size), 4)
# tmp_lr = 1e-6 + (base_lr-1e-6) * (iter_i+epoch*epoch_size) / (epoch_size * (args.wp_epoch))
set_lr(optimizer, tmp_lr)
elif epoch == args.wp_epoch and iter_i == 0:
tmp_lr = base_lr
set_lr(optimizer, tmp_lr)
# to device
images = images.to(device)
# multi-scale trick
if iter_i % 10 == 0 and iter_i > 0 and args.multi_scale:
# randomly choose a new size
size = random.randint(10, 19) * 32
input_size = [size, size]
model.set_grid(input_size)
if args.multi_scale:
# interpolate
images = torch.nn.functional.interpolate(images,
size=input_size,
mode='bilinear',
align_corners=False)
# make labels
targets = [label.tolist() for label in targets]
targets = tools.gt_creator(input_size=input_size,
stride=yolo_net.stride,
label_lists=targets)
targets = torch.tensor(targets).float().to(device)
# forward and loss
conf_loss, cls_loss, txtytwth_loss, total_loss = model(
images, target=targets)
# backprop
total_loss.backward()
optimizer.step()
optimizer.zero_grad()
if args.tfboard:
# viz loss
writer.add_scalar('object loss', conf_loss.item(),
iter_i + epoch * epoch_size)
writer.add_scalar('class loss', cls_loss.item(),
iter_i + epoch * epoch_size)
writer.add_scalar('local loss', txtytwth_loss.item(),
iter_i + epoch * epoch_size)
if iter_i % 10 == 0:
t1 = time.time()
print(
'[Epoch %d/%d][Iter %d/%d][lr %.6f]'
'[Loss: obj %.2f || cls %.2f || bbox %.2f || total %.2f || size %d || time: %.2f]'
% (epoch + 1, max_epoch, iter_i, epoch_size, tmp_lr,
conf_loss.item(), cls_loss.item(), txtytwth_loss.item(),
total_loss.item(), input_size[0], t1 - t0),
flush=True)
t0 = time.time()
if (epoch + 1) % 10 == 0:
print('Saving state, epoch:', epoch + 1)
torch.save(
model.state_dict(),
os.path.join(path_to_save,
args.version + '_' + repr(epoch + 1) + '.pth'))
# COCO evaluation
if (epoch + 1) % args.eval_epoch == 0:
model.trainable = False
model.set_grid(cfg['min_dim'])
# evaluate
ap50_95, ap50 = evaluator.evaluate(model)
print('ap50 : ', ap50)
print('ap50_95 : ', ap50_95)
# convert to training mode.
model.trainable = True
model.train()
if args.tfboard:
writer.add_scalar('val/COCOAP50', ap50, epoch + 1)
writer.add_scalar('val/COCOAP50_95', ap50_95, epoch + 1)
def train(net, device):
global cfg, hr
# set GPU
use_focal = False
if args.use_focal == 1:
print("Let's use focal loss for objectness !!!")
use_focal = True
if args.multi_scale == 1:
print('Let us use the multi-scale trick.')
ms_inds = range(len(cfg['multi_scale']))
dataset = COCODataset(data_dir=data_dir,
img_size=608,
transform=SSDAugmentation([608, 608], MEANS),
debug=args.debug)
else:
dataset = COCODataset(data_dir=data_dir,
img_size=cfg['min_dim'][0],
transform=SSDAugmentation(cfg['min_dim'], MEANS),
debug=args.debug)
print("Setting Arguments.. : ", args)
print("----------------------------------------------------------")
print('Loading the MSCOCO dataset...')
print('Training model on:', dataset.name)
print('The dataset size:', len(dataset))
print('The obj weight : ', args.obj)
print('The noobj weight : ', args.noobj)
print("----------------------------------------------------------")
input_size = cfg['min_dim']
num_classes = args.num_classes
batch_size = args.batch_size
save_folder = args.save_folder
if not os.path.exists(save_folder):
os.mkdir(save_folder)
# using tfboard
from tensorboardX import SummaryWriter
c_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time()))
log_path = 'log/yolo_v2/coco/' + c_time
if not os.path.exists(log_path):
os.mkdir(log_path)
if not os.path.exists(log_path):
os.mkdir(log_path)
writer = SummaryWriter(log_path)
if args.high_resolution == 1:
hr = True
print('Let us train yolo-v2 on the MSCOCO dataset ......')
model = net
model.to(device).train()
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
shuffle=True,
collate_fn=detection_collate,
num_workers=args.n_cpu)
evaluator = COCOAPIEvaluator(data_dir=data_dir,
img_size=cfg['min_dim'],
device=device,
transform=BaseTransform(
cfg['min_dim'], MEANS))
# optimizer setup
# optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum,
# dampening=0, weight_decay=args.weight_decay)
# optimizer = optim.Adam(model.parameters())
lr = args.lr
optimizer = optim.RMSprop(model.parameters(), lr=args.lr)
step_index = 0
epoch_size = len(dataset) // args.batch_size
# each part of loss weight
obj_w = 1.0
cla_w = 1.0
box_w = 2.0
# start training loop
iteration = 0
for epoch in range(cfg['max_epoch']):
batch_iterator = iter(dataloader)
# No WarmUp strategy or WarmUp tage has finished.
if epoch in cfg['lr_epoch']:
step_index += 1
lr = adjust_learning_rate(optimizer, args.gamma, step_index)
# COCO evaluation
if (epoch + 1) % args.eval_epoch == 0:
model.trainable = False
ap50_95, ap50 = evaluator.evaluate(model)
print('ap50 : ', ap50)
print('ap90_95 : ', ap50_95)
model.trainable = True
model.train()
writer.add_scalar('val/COCOAP50', ap50, epoch + 1)
writer.add_scalar('val/COCOAP50_95', ap50_95, epoch + 1)
# subdivision loop
optimizer.zero_grad()
for images, targets in batch_iterator:
iteration += 1
# multi-scale trick
if iteration % 10 == 0 and args.multi_scale == 1:
ms_ind = random.sample(ms_inds, 1)[0]
input_size = cfg['multi_scale'][int(ms_ind)]
# multi scale
if args.multi_scale == 1:
images = torch.nn.functional.interpolate(images,
size=input_size,
mode='bilinear',
align_corners=True)
targets = [label.tolist() for label in targets]
if args.version == 'yolo_v2':
targets = tools.gt_creator(input_size,
yolo_net.stride,
args.num_classes,
targets,
name='COCO')
elif args.version == 'yolo_v3':
targets = tools.multi_gt_creator(input_size,
yolo_net.stride,
args.num_classes,
targets,
name='COCO')
targets = torch.tensor(targets).float().to(device)
t0 = time.time()
out = model(images.to(device))
obj_loss, class_loss, box_loss = tools.loss(
out,
targets,
num_classes=args.num_classes,
use_focal=use_focal,
obj=args.obj,
noobj=args.noobj)
total_loss = obj_w * obj_loss + cla_w * class_loss + box_w * box_loss
# viz loss
writer.add_scalar('object loss', obj_loss.item(), iteration)
writer.add_scalar('class loss', class_loss.item(), iteration)
writer.add_scalar('local loss', box_loss.item(), iteration)
writer.add_scalar('total loss', total_loss.item(), iteration)
# backprop
total_loss.backward()
optimizer.step()
t1 = time.time()
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
# print(obj_loss.item(), class_loss.item(), box_loss.item())
print('Epoch[%d / %d]' % (epoch+1, cfg['max_epoch']) + ' || iter ' + repr(iteration) + \
' || Loss: %.4f ||' % (total_loss.item()) + ' || lr: %.8f ||' % (lr) + ' || input size: %d ||' % input_size[0], end=' ')
if (epoch + 1) % 10 == 0:
print('Saving state, epoch:', epoch + 1)
torch.save(
yolo_net.state_dict(), save_folder + '/' + args.version + '_' +
repr(epoch + 1) + '.pth')
def main():
"""
YOLOv3 trainer. See README for details.
"""
args = parse_args()
print("Setting Arguments.. : ", args)
cuda = torch.cuda.is_available() and args.use_cuda
os.makedirs(args.checkpoint_dir, exist_ok=True)
# Parse config settings
with open(args.cfg, 'r') as f:
cfg = yaml.load(f)
print("successfully loaded config file: ", cfg)
lr = cfg['TRAIN']['LR']
momentum = cfg['TRAIN']['MOMENTUM']
decay = cfg['TRAIN']['DECAY']
burn_in = cfg['TRAIN']['BURN_IN']
iter_size = cfg['TRAIN']['MAXITER']
steps = eval(cfg['TRAIN']['STEPS'])
batch_size = cfg['TRAIN']['BATCHSIZE']
subdivision = cfg['TRAIN']['SUBDIVISION']
ignore_thre = cfg['TRAIN']['IGNORETHRE']
random_resize = cfg['AUGMENTATION']['RANDRESIZE']
print('effective_batch_size = batch_size * iter_size = %d * %d' %
(batch_size, subdivision))
# Learning rate setup
base_lr = lr
# Initiate model
model = YOLOv3(cfg['MODEL'], ignore_thre=ignore_thre)
if args.weights_path:
print("loading darknet weights....", args.weights_path)
#parse_yolo_weights(model, args.weights_path)
model.load_weights(args.weights_path)
elif args.checkpoint:
print("loading pytorch ckpt...", args.checkpoint)
state = torch.load(args.checkpoint)
if 'model_state_dict' in state.keys():
model.load_state_dict(state['model_state_dict'])
else:
model.load_state_dict(state)
if cuda:
print("using cuda")
model = model.cuda()
#model.save_weights('../darknet/torch-yolo-tiny.weights')
if args.tfboard:
print("using tfboard")
from tensorboardX import SummaryWriter
tblogger = SummaryWriter(args.tfboard)
model.train()
imgsize = cfg['TRAIN']['IMGSIZE']
evaluator = COCOAPIEvaluator(model_type=cfg['MODEL']['TYPE'],
data_dir='COCO/',
img_size=cfg['TEST']['IMGSIZE'],
confthre=cfg['TEST']['CONFTHRE'],
nmsthre=cfg['TEST']['NMSTHRE'])
print('use image size:', cfg['TEST']['IMGSIZE'])
# optimizer setup
# set weight decay only on conv.weight
params_dict = dict(model.named_parameters())
params = []
for key, value in params_dict.items():
if 'conv.weight' in key:
params += [{
'params': value,
'weight_decay': decay * batch_size * subdivision
}]
else:
params += [{'params': value, 'weight_decay': 0.0}]
optimizer = optim.SGD(params,
lr=base_lr,
momentum=momentum,
dampening=0,
weight_decay=decay * batch_size * subdivision)
if args.checkpoint:
if 'optimizer_state_dict' in state.keys():
optimizer.load_state_dict(state['optimizer_state_dict'])
iter_state = state['iter'] + 1
# TODO: replace the following scheduler with the PyTorch's official one
ap50_95, ap50 = evaluator.evaluate(model)
def train():
args = parse_args()
data_dir = args.dataset_root
path_to_save = os.path.join(args.save_folder, args.version)
os.makedirs(path_to_save, exist_ok=True)
cfg = coco_cfg
if args.cuda:
print('use cuda')
cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
input_size = cfg['min_dim']
dataset = COCODataset(
data_dir=data_dir,
img_size=cfg['min_dim'],
transform=SSDAugmentation(cfg['min_dim'], mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229)),
debug=args.debug)
# build model
if args.version == 'centernet':
from models.centernet import CenterNet
net = CenterNet(device, input_size=input_size, num_classes=args.num_classes, trainable=True)
print('Let us train centernet on the COCO dataset ......')
else:
print('Unknown version !!!')
exit()
print("Setting Arguments.. : ", args)
print("----------------------------------------------------------")
print('Loading the MSCOCO dataset...')
print('Training model on:', dataset.name)
print('The dataset size:', len(dataset))
print("----------------------------------------------------------")
# use tfboard
if args.tfboard:
print('use tensorboard')
from torch.utils.tensorboard import SummaryWriter
c_time = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
log_path = os.path.join('log/coco/', args.version, c_time)
os.makedirs(log_path, exist_ok=True)
writer = SummaryWriter(log_path)
model = net
model.to(device).train()
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
shuffle=True,
collate_fn=detection_collate,
num_workers=args.num_workers)
evaluator = COCOAPIEvaluator(
data_dir=data_dir,
img_size=cfg['min_dim'],
device=device,
transform=BaseTransform(cfg['min_dim'], mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229))
)
# optimizer setup
base_lr = args.lr
tmp_lr = base_lr
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum,
weight_decay=args.weight_decay)
max_epoch = cfg['max_epoch']
epoch_size = len(dataset) // args.batch_size
# start training loop
t0 = time.time()
for epoch in range(max_epoch):
# use cos lr
if args.cos and epoch > 20 and epoch <= max_epoch - 20:
# use cos lr
tmp_lr = 0.00001 + 0.5*(base_lr-0.00001)*(1+math.cos(math.pi*(epoch-20)*1./ (max_epoch-20)))
set_lr(optimizer, tmp_lr)
elif args.cos and epoch > max_epoch - 20:
tmp_lr = 0.00001
set_lr(optimizer, tmp_lr)
# use step lr
else:
if epoch in cfg['lr_epoch']:
tmp_lr = tmp_lr * 0.1
set_lr(optimizer, tmp_lr)
for iter_i, (images, targets) in enumerate(dataloader):
# WarmUp strategy for learning rate
if not args.no_warm_up:
if epoch < args.wp_epoch:
tmp_lr = base_lr * pow((iter_i+epoch*epoch_size)*1. / (args.wp_epoch*epoch_size), 4)
# tmp_lr = 1e-6 + (base_lr-1e-6) * (iter_i+epoch*epoch_size) / (epoch_size * (args.wp_epoch))
set_lr(optimizer, tmp_lr)
elif epoch == args.wp_epoch and iter_i == 0:
tmp_lr = base_lr
set_lr(optimizer, tmp_lr)
targets = [label.tolist() for label in targets]
targets = tools.gt_creator(input_size, net.stride, args.num_classes, targets)
# to device
images = images.to(device)
targets = torch.tensor(targets).float().to(device)
# forward and loss
cls_loss, txty_loss, twth_loss, total_loss = model(images, target=targets)
# backprop and update
total_loss.backward()
optimizer.step()
optimizer.zero_grad()
if iter_i % 10 == 0:
if args.tfboard:
# viz loss
writer.add_scalar('class loss', cls_loss.item(), iter_i + epoch * epoch_size)
writer.add_scalar('txty loss', txty_loss.item(), iter_i + epoch * epoch_size)
writer.add_scalar('twth loss', twth_loss.item(), iter_i + epoch * epoch_size)
writer.add_scalar('total loss', total_loss.item(), iter_i + epoch * epoch_size)
t1 = time.time()
print('[Epoch %d/%d][Iter %d/%d][lr %.6f]'
'[Loss: cls %.2f || txty %.2f || twth %.2f ||total %.2f || size %d || time: %.2f]'
% (epoch+1, max_epoch, iter_i, epoch_size, tmp_lr,
cls_loss.item(), txty_loss.item(), twth_loss.item(), total_loss.item(), input_size, t1-t0),
flush=True)
t0 = time.time()
# COCO evaluation
if (epoch + 1) % args.eval_epoch == 0:
model.trainable = False
# evaluate
ap50_95, ap50 = evaluator.evaluate(model)
print('ap50 : ', ap50)
print('ap50_95 : ', ap50_95)
# convert to training mode.
model.trainable = True
model.train()
if args.tfboard:
writer.add_scalar('val/COCOAP50', ap50, epoch + 1)
writer.add_scalar('val/COCOAP50_95', ap50_95, epoch + 1)
if (epoch + 1) % 10 == 0:
print('Saving state, epoch:', epoch + 1)
torch.save(model.state_dict(), os.path.join(path_to_save,
args.version + '_' + repr(epoch + 1) + '.pth')
)
def main():
"""
YOLOv3 trainer. See README for details.
"""
args = parse_args()
print("Setting Arguments.. : ", args)
cuda = torch.cuda.is_available() and args.use_cuda
os.makedirs(args.checkpoint_dir, exist_ok=True)
# Parse config settings
with open(args.cfg, 'r') as f:
cfg = yaml.load(f)
print("successfully loaded config file: ", cfg)
momentum = cfg['TRAIN']['MOMENTUM']
decay = cfg['TRAIN']['DECAY']
burn_in = cfg['TRAIN']['BURN_IN']
iter_size = cfg['TRAIN']['MAXITER']
steps = eval(cfg['TRAIN']['STEPS'])
batch_size = cfg['TRAIN']['BATCHSIZE']
subdivision = cfg['TRAIN']['SUBDIVISION']
ignore_thre = cfg['TRAIN']['IGNORETHRE']
random_resize = cfg['AUGMENTATION']['RANDRESIZE']
base_lr = cfg['TRAIN']['LR'] / batch_size / subdivision
print('effective_batch_size = batch_size * iter_size = %d * %d' %
(batch_size, subdivision))
# Learning rate setup
def burnin_schedule(i):
if i < burn_in:
factor = pow(i / burn_in, 4)
elif i < steps[0]:
factor = 1.0
elif i < steps[1]:
factor = 0.1
else:
factor = 0.01
return factor
# Initiate model
model = YOLOv3(cfg['MODEL'], ignore_thre=ignore_thre)
if args.weights_path:
print("loading darknet weights....", args.weights_path)
parse_yolo_weights(model, args.weights_path)
elif args.checkpoint:
print("loading pytorch ckpt...", args.checkpoint)
state = torch.load(args.checkpoint)
if 'model_state_dict' in state.keys():
model.load_state_dict(state['model_state_dict'])
else:
model.load_state_dict(state)
if cuda:
print("using cuda")
model = model.cuda()
if args.tfboard:
print("using tfboard")
from tensorboardX import SummaryWriter
tblogger = SummaryWriter(args.tfboard)
model.train()
imgsize = cfg['TRAIN']['IMGSIZE']
dataset = COCODataset(model_type=cfg['MODEL']['TYPE'],
data_dir='COCO/',
img_size=imgsize,
augmentation=cfg['AUGMENTATION'],
debug=args.debug)
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=args.n_cpu)
dataiterator = iter(dataloader)
evaluator = COCOAPIEvaluator(model_type=cfg['MODEL']['TYPE'],
data_dir='COCO/',
img_size=cfg['TEST']['IMGSIZE'],
confthre=cfg['TEST']['CONFTHRE'],
nmsthre=cfg['TEST']['NMSTHRE'])
dtype = torch.cuda.FloatTensor if cuda else torch.FloatTensor
# optimizer setup
# set weight decay only on conv.weight
params_dict = dict(model.named_parameters())
params = []
for key, value in params_dict.items():
if 'conv.weight' in key:
params += [{'params':value, 'weight_decay':decay * batch_size * subdivision}]
else:
params += [{'params':value, 'weight_decay':0.0}]
optimizer = optim.SGD(params, lr=base_lr, momentum=momentum,
dampening=0, weight_decay=decay * batch_size * subdivision)
iter_state = 0
if args.checkpoint:
if 'optimizer_state_dict' in state.keys():
optimizer.load_state_dict(state['optimizer_state_dict'])
iter_state = state['iter'] + 1
scheduler = optim.lr_scheduler.LambdaLR(optimizer, burnin_schedule)
# start training loop
for iter_i in range(iter_state, iter_size + 1):
# COCO evaluation
if iter_i % args.eval_interval == 0 and iter_i > 0:
ap50_95, ap50 = evaluator.evaluate(model)
model.train()
if args.tfboard:
tblogger.add_scalar('val/COCOAP50', ap50, iter_i)
tblogger.add_scalar('val/COCOAP50_95', ap50_95, iter_i)
# subdivision loop
optimizer.zero_grad()
for inner_iter_i in range(subdivision):
try:
imgs, targets, _, _ = next(dataiterator) # load a batch
except StopIteration:
dataiterator = iter(dataloader)
imgs, targets, _, _ = next(dataiterator) # load a batch
imgs = Variable(imgs.type(dtype))
targets = Variable(targets.type(dtype), requires_grad=False)
loss = model(imgs, targets)
loss.backward()
optimizer.step()
scheduler.step()
if iter_i % 10 == 0:
# logging
current_lr = scheduler.get_lr()[0] * batch_size * subdivision
print('[Iter %d/%d] [lr %f] '
'[Losses: xy %f, wh %f, conf %f, cls %f, total %f, imgsize %d]'
% (iter_i, iter_size, current_lr,
model.loss_dict['xy'], model.loss_dict['wh'],
model.loss_dict['conf'], model.loss_dict['cls'],
model.loss_dict['l2'], imgsize),
flush=True)
if args.tfboard:
tblogger.add_scalar('train/total_loss', model.loss_dict['l2'], iter_i)
# random resizing
if random_resize:
imgsize = (random.randint(0, 9) % 10 + 10) * 32
dataset.img_shape = (imgsize, imgsize)
dataset.img_size = imgsize
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=args.n_cpu)
dataiterator = iter(dataloader)
# save checkpoint
if iter_i > 0 and (iter_i % args.checkpoint_interval == 0):
torch.save({'iter': iter_i,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
},
os.path.join(args.checkpoint_dir, "snapshot"+str(iter_i)+".ckpt"))
if args.tfboard:
tblogger.close()
def train():
args = parse_args()
path_to_save = os.path.join(args.save_folder, args.dataset, args.version)
os.makedirs(path_to_save, exist_ok=True)
# cuda
if args.cuda:
print('use cuda')
cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
# mosaic augmentation
if args.mosaic:
print('use Mosaic Augmentation ...')
# multi-scale
if args.multi_scale:
print('use the multi-scale trick ...')
train_size = [640, 640]
val_size = [512, 512]
else:
train_size = [512, 512]
val_size = [512, 512]
cfg = train_cfg
# dataset and evaluator
print("Setting Arguments.. : ", args)
print("----------------------------------------------------------")
print('Loading the dataset...')
if args.dataset == 'voc':
data_dir = VOC_ROOT
num_classes = 20
dataset = VOCDetection(root=data_dir, img_size=train_size[0],
transform=SSDAugmentation(train_size),
mosaic=args.mosaic
)
evaluator = VOCAPIEvaluator(data_root=data_dir,
img_size=val_size,
device=device,
transform=BaseTransform(val_size),
labelmap=VOC_CLASSES
)
elif args.dataset == 'coco':
data_dir = coco_root
num_classes = 80
dataset = COCODataset(
data_dir=data_dir,
img_size=train_size[0],
transform=SSDAugmentation(train_size),
debug=args.debug,
mosaic=args.mosaic
)
evaluator = COCOAPIEvaluator(
data_dir=data_dir,
img_size=val_size,
device=device,
transform=BaseTransform(val_size)
)
else:
print('unknow dataset !! Only support voc and coco !!')
exit(0)
print('Training model on:', dataset.name)
print('The dataset size:', len(dataset))
print("----------------------------------------------------------")
# dataloader
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
shuffle=True,
collate_fn=detection_collate,
num_workers=args.num_workers,
pin_memory=True
)
# build model
if args.version == 'centernet':
from models.centernet import CenterNet
net = CenterNet(device, input_size=train_size, num_classes=num_classes, trainable=True)
print('Let us train centernet on the %s dataset ......' % (args.dataset))
else:
print('Unknown version !!!')
exit()
model = net
model.to(device).train()
# use tfboard
if args.tfboard:
print('use tensorboard')
from torch.utils.tensorboard import SummaryWriter
c_time = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
log_path = os.path.join('log/coco/', args.version, c_time)
os.makedirs(log_path, exist_ok=True)
writer = SummaryWriter(log_path)
# keep training
if args.resume is not None:
print('keep training model: %s' % (args.resume))
model.load_state_dict(torch.load(args.resume, map_location=device))
# optimizer setup
base_lr = args.lr
tmp_lr = base_lr
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay
)
max_epoch = cfg['max_epoch']
epoch_size = len(dataset) // args.batch_size
# start training loop
t0 = time.time()
for epoch in range(args.start_epoch, max_epoch):
# use cos lr
if args.cos and epoch > 20 and epoch <= max_epoch - 20:
# use cos lr
tmp_lr = 0.00001 + 0.5*(base_lr-0.00001)*(1+math.cos(math.pi*(epoch-20)*1./ (max_epoch-20)))
set_lr(optimizer, tmp_lr)
elif args.cos and epoch > max_epoch - 20:
tmp_lr = 0.00001
set_lr(optimizer, tmp_lr)
# use step lr
else:
if epoch in cfg['lr_epoch']:
tmp_lr = tmp_lr * 0.1
set_lr(optimizer, tmp_lr)
for iter_i, (images, targets) in enumerate(dataloader):
# WarmUp strategy for learning rate
if not args.no_warm_up:
if epoch < args.wp_epoch:
tmp_lr = base_lr * pow((iter_i+epoch*epoch_size)*1. / (args.wp_epoch*epoch_size), 4)
# tmp_lr = 1e-6 + (base_lr-1e-6) * (iter_i+epoch*epoch_size) / (epoch_size * (args.wp_epoch))
set_lr(optimizer, tmp_lr)
elif epoch == args.wp_epoch and iter_i == 0:
tmp_lr = base_lr
set_lr(optimizer, tmp_lr)
# to device
images = images.to(device)
# multi-scale trick
if iter_i % 10 == 0 and iter_i > 0 and args.multi_scale:
# randomly choose a new size
size = random.randint(10, 19) * 32
train_size = [size, size]
model.set_grid(train_size)
if args.multi_scale:
# interpolate
images = torch.nn.functional.interpolate(images, size=train_size, mode='bilinear', align_corners=False)
# make train label
targets = [label.tolist() for label in targets]
targets = tools.gt_creator(train_size, net.stride, args.num_classes, targets)
targets = torch.tensor(targets).float().to(device)
# forward and loss
cls_loss, txty_loss, twth_loss, total_loss = model(images, target=targets)
# backprop
total_loss.backward()
optimizer.step()
optimizer.zero_grad()
if iter_i % 10 == 0:
if args.tfboard:
# viz loss
writer.add_scalar('class loss', cls_loss.item(), iter_i + epoch * epoch_size)
writer.add_scalar('txty loss', txty_loss.item(), iter_i + epoch * epoch_size)
writer.add_scalar('twth loss', twth_loss.item(), iter_i + epoch * epoch_size)
writer.add_scalar('total loss', total_loss.item(), iter_i + epoch * epoch_size)
t1 = time.time()
print('[Epoch %d/%d][Iter %d/%d][lr %.6f]'
'[Loss: cls %.2f || txty %.2f || twth %.2f ||total %.2f || size %d || time: %.2f]'
% (epoch+1, max_epoch, iter_i, epoch_size, tmp_lr,
cls_loss.item(), txty_loss.item(), twth_loss.item(), total_loss.item(), train_size[0], t1-t0),
flush=True)
t0 = time.time()
# evaluation
if (epoch) % args.eval_epoch == 0:
model.trainable = False
model.set_grid(val_size)
model.eval()
# evaluate
evaluator.evaluate(model)
# convert to training mode.
model.trainable = True
model.set_grid(train_size)
model.train()
# save model
if (epoch + 1) % 10 == 0:
print('Saving state, epoch:', epoch + 1)
torch.save(model.state_dict(), os.path.join(path_to_save,
args.version + '_' + repr(epoch + 1) + '.pth')
)
def main():
args = parse()
np.random.seed(args.seed)
print('arguments: ', args)
# Model setup
if args.dataset == 'coco2017':
train_data = COCODataset()
test_data = COCODataset(json_file='instances_val2017.json',
name='val2017',
id_list_file='val2017.txt')
train_class_ids = train_data.class_ids
test_ids = test_data.ids
cocoanns = test_data.coco
if args.extractor == 'vgg16':
mask_rcnn = MaskRCNNVGG16(n_fg_class=80,
pretrained_model=args.pretrained,
roi_size=args.roi_size,
roi_align=args.roialign)
elif args.extractor == 'resnet50':
mask_rcnn = MaskRCNNResNet(n_fg_class=80,
pretrained_model=args.pretrained,
roi_size=args.roi_size,
n_layers=50,
roi_align=args.roialign,
class_ids=train_class_ids)
elif args.extractor == 'resnet101':
mask_rcnn = MaskRCNNResNet(n_fg_class=80,
pretrained_model=args.pretrained,
roi_size=args.roi_size,
n_layers=101,
roi_align=args.roialign,
class_ids=train_class_ids)
mask_rcnn.use_preset('evaluate')
model = MaskRCNNTrainChain(mask_rcnn,
gamma=args.gamma,
roi_size=args.roi_size)
# Trainer setup
if args.gpu >= 0:
chainer.cuda.get_device_from_id(args.gpu).use()
model.to_gpu()
optimizer = chainer.optimizers.MomentumSGD(lr=args.lr, momentum=0.9)
#optimizer = chainer.optimizers.Adam()#alpha=0.001, beta1=0.9, beta2=0.999 , eps=0.00000001)
optimizer.setup(model)
optimizer.add_hook(chainer.optimizer.WeightDecay(rate=0.0001))
train_data = TransformDataset(train_data,
Transform(mask_rcnn, train_class_ids))
test_data = TransformDataset(test_data,
Transform(mask_rcnn, train_class_ids))
train_iter = chainer.iterators.SerialIterator(train_data,
batch_size=args.batchsize)
test_iter = chainer.iterators.SerialIterator(test_data,
batch_size=1,
repeat=False,
shuffle=False)
updater = SubDivisionUpdater(train_iter,
optimizer,
device=args.gpu,
subdivisions=args.batchsize)
#updater = ParallelUpdater(train_iter, optimizer, devices={"main": 0, "second": 1}, converter=convert ) #for training with multiple GPUs
trainer = training.Trainer(updater, (args.iteration, 'iteration'),
out=args.out)
# Extensions
trainer.extend(extensions.snapshot_object(model.mask_rcnn,
'snapshot_model.npz'),
trigger=(args.snapshot, 'iteration'))
trainer.extend(extensions.ExponentialShift('lr', 10),
trigger=ManualScheduleTrigger([args.lr_initialchange],
'iteration'))
trainer.extend(extensions.ExponentialShift('lr', 0.1),
trigger=(args.lr_step, 'iteration'))
if args.resume is not None:
chainer.serializers.load_npz(args.resume, model.mask_rcnn)
if args.freeze_bn:
freeze_bn(model.mask_rcnn)
if args.bn2affine:
bn_to_affine(model.mask_rcnn)
log_interval = 40, 'iteration'
plot_interval = 160, 'iteration'
print_interval = 40, 'iteration'
#trainer.extend(extensions.Evaluator(test_iter, model, device=args.gpu), trigger=(args.validation, 'iteration'))
#trainer.extend(DetectionCOCOEvaluator(test_iter, model.mask_rcnn), trigger=(args.validation, 'iteration')) #COCO AP Evaluator with VOC metric
trainer.extend(COCOAPIEvaluator(test_iter, model.mask_rcnn, test_ids,
cocoanns),
trigger=(args.validation, 'iteration')) #COCO AP Evaluator
trainer.extend(chainer.training.extensions.observe_lr(),
trigger=log_interval)
trainer.extend(extensions.LogReport(trigger=log_interval))
trainer.extend(extensions.PrintReport([
'iteration',
'epoch',
'elapsed_time',
'lr',
'main/loss',
'main/avg_loss',
'main/roi_loc_loss',
'main/roi_cls_loss',
'main/roi_mask_loss',
'main/rpn_loc_loss',
'main/rpn_cls_loss',
'validation/main/loss',
'validation/main/map',
]),
trigger=print_interval)
trainer.extend(extensions.ProgressBar(update_interval=1000))
#trainer.extend(extensions.dump_graph('main/loss'))
try:
trainer.run()
except:
traceback.print_exc()
def eval():
"""
YOLOv3 evaler. See README for details.
"""
args = parse_args()
print("Setting Arguments.. : ", args)
cuda = torch.cuda.is_available() and args.use_cuda
if args.distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend="nccl", init_method="env://")
# Parse config settings
with open(args.cfg, 'r') as f:
cfg = yaml.safe_load(f)
print("successfully loaded config file: ", cfg)
test_size = (args.test_size,args.test_size)
if args.dataset == 'COCO':
evaluator = COCOAPIEvaluator(
data_dir='data/COCO/',
img_size=test_size,
confthre=0.001,
nmsthre=0.65,
testset=args.testset,
vis=args.vis)
num_class=80
elif args.dataset == 'VOC':
'''
# COCO style evaluation, you have to convert xml annotation files into a json file.
evaluator = COCOAPIEvaluator(
data_dir='data/VOC/',
img_size=test_size,
confthre=cfg['TEST']['CONFTHRE'],
nmsthre=cfg['TEST']['NMSTHRE'],
testset=args.testset,
voc = True)
'''
evaluator = VOCEvaluator(
data_dir='data/VOC/',
img_size=test_size,
confthre=0.001,
nmsthre=0.65,
vis=args.vis)
num_class=20
# Initiate model
if args.asff:
from models.yolov3_asff import YOLOv3
print('Testing YOLOv3 with ASFF!')
model = YOLOv3(num_classes = num_class, rfb=args.rfb, vis=args.vis)
else:
from models.yolov3_baseline import YOLOv3
print('Testing YOLOv3 strong baseline!')
if args.vis:
print('Visualization is not supported for YOLOv3 baseline model')
args.vis = False
model = YOLOv3(num_classes = num_class, rfb=args.rfb)
save_to_disk = (not args.distributed) or distributed_util.get_rank() == 0
if args.checkpoint:
print("loading pytorch ckpt...", args.checkpoint)
cpu_device = torch.device("cpu")
ckpt = torch.load(args.checkpoint, map_location=cpu_device)
model.load_state_dict(ckpt,strict=False)
if cuda:
print("using cuda")
torch.backends.cudnn.benchmark = True
device = torch.device("cuda")
model = model.to(device)
if args.half:
model = model.half()
if args.ngpu > 1:
if args.distributed:
model = apex.parallel.DistributedDataParallel(model, delay_allreduce=True)
#model = apex.parallel.DistributedDataParallel(model)
else:
model = nn.DataParallel(model)
dtype = torch.float16 if args.half else torch.float32
if args.distributed:
distributed_util.synchronize()
ap50_95, ap50 = evaluator.evaluate(model, args.half, args.distributed)
if args.distributed:
distributed_util.synchronize()
sys.exit(0)
def main():
"""
YOLOv3 trainer. See README for details.
"""
date = time.strftime('%Y-%m-%d-%H-%M', time.localtime())
args = parse_args()
print("Setting Arguments.. : ", args)
cuda = torch.cuda.is_available() and args.use_cuda
# Parse config settings
with open(args.cfg, 'r') as f:
cfg = yaml.load(f)
print("successfully loaded config file: ", cfg)
lr = cfg['TRAIN']['LR']
momentum = cfg['TRAIN']['MOMENTUM']
decay = cfg['TRAIN']['DECAY']
burn_in = cfg['TRAIN']['BURN_IN']
iter_size = cfg['TRAIN']['MAXITER']
steps = eval(cfg['TRAIN']['STEPS'])
batch_size = cfg['TRAIN']['BATCHSIZE']
subdivision = cfg['TRAIN']['SUBDIVISION']
ignore_thre = cfg['TRAIN']['IGNORETHRE']
random_resize = cfg['AUGMENTATION']['RANDRESIZE']
print('effective_batch_size = batch_size * iter_size = %d * %d' %
(batch_size, subdivision))
# Learning rate setup
base_lr = lr
# Initiate model
model = YOLOv3(cfg['MODEL'], ignore_thre=ignore_thre)
if args.pretrained:
if cfg['MODEL']['TYPE'] == 'YOLOv3':
# self.module_list = create_yolov3_modules(config_model, ignore_thre)
print('load yolov3 pretrained model')
model.load_pretrained_weights('weights/darknet53.conv.74', cutoff=75)
elif cfg['MODEL']['TYPE'].lower() == 'yolov3-tiny':
print('load yolov3 pretrained model')
model.load_pretrained_weights('weights/yolov3-tiny.conv.15', cutoff=16)
else:
print('no this type pretrained model')
if args.weights_path:
print("loading darknet weights....", args.weights_path)
#parse_yolo_weights(model, args.weights_path)
model.load_weights(args.weights_path)
elif args.checkpoint:
print("loading pytorch ckpt...", args.checkpoint)
state = torch.load(args.checkpoint)
if 'model_state_dict' in state.keys():
model.load_state_dict(state['model_state_dict'])
else:
model.load_state_dict(state)
if cuda:
print("using cuda")
model = model.cuda()
#model.save_weights('../darknet/torch-yolo-tiny.weights')
if args.tfboard:
print("using tfboard")
from tensorboardX import SummaryWriter
tblogger = SummaryWriter(args.tfboard)
#model = torch.nn.DataParallel(model)
model.train()
imgsize = cfg['TRAIN']['IMGSIZE']
dataset = COCODataset(model_type=cfg['MODEL']['TYPE'],
data_dir='COCO/',
img_size=imgsize,
augmentation=cfg['AUGMENTATION'],
debug=args.debug)
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=args.n_cpu)
dataiterator = iter(dataloader)
evaluator = COCOAPIEvaluator(model_type=cfg['MODEL']['TYPE'],
data_dir='COCO/',
img_size=cfg['TEST']['IMGSIZE'],
confthre=cfg['TEST']['CONFTHRE'],
nmsthre=cfg['TEST']['NMSTHRE'])
dtype = torch.cuda.FloatTensor if cuda else torch.FloatTensor
# optimizer setup
# set weight decay only on conv.weight
params_dict = dict(model.named_parameters())
params = []
for key, value in params_dict.items():
if 'conv.weight' in key:
params += [{'params':value, 'weight_decay':decay * batch_size * subdivision}]
else:
params += [{'params':value, 'weight_decay':0.0}]
optimizer = optim.SGD(params, lr=base_lr, momentum=momentum,
dampening=0, weight_decay=decay * batch_size * subdivision)
iter_state = 0
if args.checkpoint:
if 'optimizer_state_dict' in state.keys():
optimizer.load_state_dict(state['optimizer_state_dict'])
iter_state = state['iter'] + 1
# TODO: replace the following scheduler with the PyTorch's official one
tmp_lr = base_lr
def set_lr(tmp_lr):
#print('set lr:', tmp_lr / batch_size / subdivision)
for param_group in optimizer.param_groups:
param_group['lr'] = tmp_lr / batch_size / subdivision
#param_group['lr'] = tmp_lr / subdivision
# start training loop
for iter_i in range(iter_state, iter_size + 1):
# COCO evaluation
if iter_i % args.eval_interval == 0 and iter_i > 0:
ap50_95, ap50 = evaluator.evaluate(model)
model.train()
if args.tfboard:
tblogger.add_scalar('val/COCOAP50', ap50, iter_i)
tblogger.add_scalar('val/COCOAP50_95', ap50_95, iter_i)
# learning rate scheduling
if iter_i < burn_in:
tmp_lr = base_lr * pow(iter_i / burn_in, 4)
set_lr(tmp_lr)
elif iter_i == burn_in:
tmp_lr = base_lr
set_lr(tmp_lr)
elif iter_i in steps:
tmp_lr = tmp_lr * 0.1
set_lr(tmp_lr)
# subdivision loop
optimizer.zero_grad()
for inner_iter_i in range(subdivision):
try:
imgs, targets, _, _ = next(dataiterator) # load a batch
except StopIteration:
dataiterator = iter(dataloader)
imgs, targets, _, _ = next(dataiterator) # load a batch
imgs = Variable(imgs.type(dtype))
targets = Variable(targets.type(dtype), requires_grad=False)
loss = model(imgs, targets)
loss.backward()
optimizer.step()
if iter_i % 10 == 0:
# logging
total_loss = model.loss_dict['xy'] + model.loss_dict['wh'] + model.loss_dict['conf'] + model.loss_dict['cls']
print('[Iter %d/%d] [lr %f] '
'[Losses: xy %f, wh %f, conf %f, cls %f, l2 %f, total %f, imgsize %d]'
% (iter_i, iter_size, tmp_lr,
model.loss_dict['xy'], model.loss_dict['wh'],
model.loss_dict['conf'], model.loss_dict['cls'],
model.loss_dict['l2'], total_loss , imgsize),
flush=True)
if args.tfboard:
tblogger.add_scalar('train/total_loss', model.loss_dict['l2'], iter_i)
# random resizing
if random_resize:
imgsize = (random.randint(0, 9) % 10 + 10) * 32
dataset.img_shape = (imgsize, imgsize)
dataset.img_size = imgsize
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=args.n_cpu)
dataiterator = iter(dataloader)
# save checkpoint
if iter_i > 0 and (iter_i % args.checkpoint_interval == 0):
try:
args.checkpoint_dir = os.path.join(args.checkpoint_dir, args.snap_name + '_' + date)
#print(args.checkpoint_dir)
os.makedirs(args.checkpoint_dir, exist_ok=True)
time.sleep(5)
torch.save({'iter': iter_i,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
},
os.path.join(args.checkpoint_dir, "snapshot"+str(iter_i)+".pth"))
time.sleep(5)
model.save_weights(os.path.join(args.checkpoint_dir, "snapshot"+str(iter_i)+".weights"))
time.sleep(5)
except Exception as e:
print(e)
if args.tfboard:
tblogger.close()
def train(net, device):
global cfg, hr
# set GPU
use_focal = False
if args.use_focal:
print("Let's use focal loss for objectness !!!")
use_focal = True
if args.multi_scale:
print('Let us use the multi-scale trick.')
ms_inds = range(len(cfg['multi_scale']))
dataset = COCODataset(
data_dir=data_dir,
img_size=608,
transform=SSDAugmentation([608, 608], mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229)),
debug=args.debug)
else:
dataset = COCODataset(
data_dir=data_dir,
img_size=cfg['min_dim'][0],
transform=SSDAugmentation(cfg['min_dim'], mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229)),
debug=args.debug)
print("Setting Arguments.. : ", args)
print("----------------------------------------------------------")
print('Loading the MSCOCO dataset...')
print('Training model on:', dataset.name)
print('The dataset size:', len(dataset))
print('The obj weight : ', args.obj)
print('The noobj weight : ', args.noobj)
print("----------------------------------------------------------")
input_size = cfg['min_dim']
num_classes = args.num_classes
batch_size = args.batch_size
os.makedirs(args.save_folder + args.version, exist_ok=True)
# using tfboard
from tensorboardX import SummaryWriter
c_time = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time()))
log_path = 'log/coco/' + c_time
os.makedirs(log_path, exist_ok=True)
writer = SummaryWriter(log_path)
if args.high_resolution == 1:
hr = True
print('Let us train yolo-v2 on the MSCOCO dataset ......')
model = net
model.to(device).train()
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
shuffle=True,
collate_fn=detection_collate,
num_workers=args.n_cpu)
evaluator = COCOAPIEvaluator(
data_dir=data_dir,
img_size=cfg['min_dim'],
device=device,
transform=BaseTransform(cfg['min_dim'], MEANS)
)
# optimizer setup
lr = args.lr
# optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum,
# weight_decay=args.weight_decay)
# optimizer = optim.Adam(model.parameters())
optimizer = optim.RMSprop(model.parameters(), lr=args.lr)
step_index = 0
epoch_size = len(dataset) // args.batch_size
# each part of loss weight
obj_w = 1.0
cla_w = 1.0
box_w = 1.0
# start training loop
iteration = 0
t0 = time.time()
for epoch in range(cfg['max_epoch']):
batch_iterator = iter(dataloader)
# use cos lr
if args.cos and epoch > 20 and epoch <= cfg['max_epoch'] - 20:
# use cos lr
lr = cos_lr(optimizer, epoch, cfg['max_epoch'])
elif args.cos and epoch > cfg['max_epoch'] - 20:
lr = 0.00001
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# use step lr
else:
if epoch in cfg['lr_epoch']:
step_index += 1
lr = adjust_learning_rate(optimizer, args.gamma, step_index)
# COCO evaluation
if (epoch + 1) % args.eval_epoch == 0:
model.trainable = False
ap50_95, ap50 = evaluator.evaluate(model)
print('ap50 : ', ap50)
print('ap50_95 : ', ap50_95)
model.trainable = True
model.train()
writer.add_scalar('val/COCOAP50', ap50, epoch + 1)
writer.add_scalar('val/COCOAP50_95', ap50_95, epoch + 1)
# subdivision loop
for images, targets in batch_iterator:
# WarmUp strategy for learning rate
if not args.no_warm_up == 'yes':
if epoch < args.wp_epoch:
lr = warmup_strategy(optimizer, epoch_size, iteration)
iteration += 1
# multi-scale trick
if iteration % 10 == 0 and args.multi_scale:
ms_ind = random.sample(ms_inds, 1)[0]
input_size = cfg['multi_scale'][int(ms_ind)]
# multi scale
if args.multi_scale:
images = torch.nn.functional.interpolate(images, size=input_size, mode='bilinear', align_corners=True)
targets = [label.tolist() for label in targets]
if args.version == 'yolo_v2' or args.version == 'tiny_yolo_v2':
targets = tools.gt_creator(input_size, yolo_net.stride, targets, name='COCO')
elif args.version == 'yolo_v3' or args.version == 'tiny_yolo_v3':
targets = tools.multi_gt_creator(input_size, yolo_net.stride, targets, name='COCO')
targets = torch.tensor(targets).float().to(device)
out = model(images.to(device))
optimizer.zero_grad()
obj_loss, class_loss, box_loss = tools.loss(out, targets, num_classes=args.num_classes,
use_focal=use_focal,
obj=args.obj,
noobj=args.noobj)
total_loss = obj_w * obj_loss + cla_w * class_loss + box_w * box_loss
# viz loss
writer.add_scalar('object loss', obj_loss.item(), iteration)
writer.add_scalar('class loss', class_loss.item(), iteration)
writer.add_scalar('local loss', box_loss.item(), iteration)
writer.add_scalar('total loss', total_loss.item(), iteration)
# backprop
total_loss.backward()
optimizer.step()
if iteration % 10 == 0:
t1 = time.time()
print('[Epoch %d/%d][Iter %d][lr %.8f]'
'[Loss: obj %.2f || cls %.2f || bbox %.2f || total %.2f || imgsize %d || time: %.2f]'
% (epoch+1, cfg['max_epoch'], iteration, lr,
obj_loss.item(), class_loss.item(), box_loss.item(), total_loss.item(), input_size[0], t1-t0),
flush=True)
t0 = time.time()
if (epoch + 1) % 10 == 0:
print('Saving state, epoch:', epoch + 1)
torch.save(yolo_net.state_dict(), os.path.join(args.save_folder + args.version,
args.version + '_' + repr(epoch + 1) + '.pth')
)
def train():
args = parse_args()
data_dir = args.dataset_root
path_to_save = os.path.join(args.save_folder, args.version)
os.makedirs(path_to_save, exist_ok=True)
hr = False
if args.high_resolution:
print('use hi-res backbone')
hr = True
cfg = coco_ab
if args.cuda:
print('use cuda')
cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
if args.multi_scale:
print('Let us use the multi-scale trick.')
input_size = [608, 608]
dataset = COCODataset(data_dir=data_dir,
img_size=608,
transform=SSDAugmentation([608, 608],
mean=(0.406, 0.456,
0.485),
std=(0.225, 0.224,
0.229)),
debug=args.debug)
else:
input_size = cfg['min_dim']
dataset = COCODataset(data_dir=data_dir,
img_size=cfg['min_dim'][0],
transform=SSDAugmentation(cfg['min_dim'],
mean=(0.406, 0.456,
0.485),
std=(0.225, 0.224,
0.229)),
debug=args.debug)
# build model
if args.version == 'yolo_v2':
from models.yolo_v2 import myYOLOv2
total_anchor_size = tools.get_total_anchor_size(name='COCO')
yolo_net = myYOLOv2(device,
input_size=input_size,
num_classes=args.num_classes,
trainable=True,
anchor_size=total_anchor_size,
hr=hr)
print('Let us train yolo-v2 on the COCO dataset ......')
elif args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
total_anchor_size = tools.get_total_anchor_size(multi_level=True,
name='COCO',
version='yolo_v3')
yolo_net = myYOLOv3(device,
input_size=input_size,
num_classes=args.num_classes,
trainable=True,
anchor_size=total_anchor_size,
hr=hr)
print('Let us train yolo-v3 on the COCO dataset ......')
elif args.version == 'slim_yolo_v2':
from models.slim_yolo_v2 import SlimYOLOv2
total_anchor_size = tools.get_total_anchor_size(name='COCO')
yolo_net = SlimYOLOv2(device,
input_size=input_size,
num_classes=args.num_classes,
trainable=True,
anchor_size=total_anchor_size,
hr=hr)
print('Let us train slim-yolo-v2 on the COCO dataset ......')
elif args.version == 'tiny_yolo_v3':
from models.tiny_yolo_v3 import YOLOv3tiny
total_anchor_size = tools.get_total_anchor_size(multi_level=True,
name='COCO',
version='tiny_yolo_v3')
yolo_net = YOLOv3tiny(device,
input_size=input_size,
num_classes=args.num_classes,
trainable=True,
anchor_size=total_anchor_size,
hr=hr)
print('Let us train tiny-yolo-v3 on the COCO dataset ......')
else:
print('Unknown version !!!')
exit()
print("Setting Arguments.. : ", args)
print("----------------------------------------------------------")
print('Loading the MSCOCO dataset...')
print('Training model on:', dataset.name)
print('The dataset size:', len(dataset))
print("----------------------------------------------------------")
# use tfboard
if args.tfboard:
print('use tensorboard')
from torch.utils.tensorboard import SummaryWriter
c_time = time.strftime('%Y-%m-%d %H:%M:%S',
time.localtime(time.time()))
log_path = os.path.join('log/coco/', args.version, c_time)
os.makedirs(log_path, exist_ok=True)
writer = SummaryWriter(log_path)
model = yolo_net
# keep training
if args.resume is not None:
print('keep training model: %s' % (args.resume))
model.load_state_dict(torch.load(args.resume, map_location=device))
model.to(device).train()
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
collate_fn=detection_collate,
num_workers=args.num_workers)
evaluator = COCOAPIEvaluator(data_dir=data_dir,
img_size=cfg['min_dim'],
device=device,
transform=BaseTransform(cfg['min_dim'],
mean=(0.406, 0.456,
0.485),
std=(0.225, 0.224,
0.229)))
# optimizer setup
base_lr = args.lr
tmp_lr = base_lr
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
max_epoch = cfg['max_epoch']
epoch_size = len(dataset) // args.batch_size
# start training loop
t0 = time.time()
for epoch in range(args.start_epoch, max_epoch):
# use cos lr
if args.cos and epoch > 20 and epoch <= max_epoch - 20:
# use cos lr
tmp_lr = 0.00001 + 0.5 * (base_lr - 0.00001) * (
1 + math.cos(math.pi * (epoch - 20) * 1. / (max_epoch - 20)))
set_lr(optimizer, tmp_lr)
elif args.cos and epoch > max_epoch - 20:
tmp_lr = 0.00001
set_lr(optimizer, tmp_lr)
# use step lr
else:
if epoch in cfg['lr_epoch']:
tmp_lr = tmp_lr * 0.1
set_lr(optimizer, tmp_lr)
for iter_i, (images, targets) in enumerate(dataloader):
# WarmUp strategy for learning rate
if not args.no_warm_up:
if epoch < args.wp_epoch:
tmp_lr = base_lr * pow((iter_i + epoch * epoch_size) * 1. /
(args.wp_epoch * epoch_size), 4)
# tmp_lr = 1e-6 + (base_lr-1e-6) * (iter_i+epoch*epoch_size) / (epoch_size * (args.wp_epoch))
set_lr(optimizer, tmp_lr)
elif epoch == args.wp_epoch and iter_i == 0:
tmp_lr = base_lr
set_lr(optimizer, tmp_lr)
targets = [label.tolist() for label in targets]
if args.version == 'yolo_v2' or args.version == 'slim_yolo_v2':
targets = tools.gt_creator(input_size,
yolo_net.stride,
targets,
name='COCO',
version=args.version)
elif args.version == 'yolo_v3' or args.version == 'tiny_yolo_v3':
targets = tools.multi_gt_creator(input_size,
yolo_net.stride,
targets,
name='COCO',
version=args.version)
# to device
images = images.to(device)
targets = torch.tensor(targets).float().to(device)
# forward and loss
conf_loss, cls_loss, txtytwth_loss, total_loss = model(
images, target=targets)
# backprop
total_loss.backward()
optimizer.step()
optimizer.zero_grad()
if iter_i % 10 == 0:
if args.tfboard:
# viz loss
writer.add_scalar('object loss', conf_loss.item(),
iter_i + epoch * epoch_size)
writer.add_scalar('class loss', cls_loss.item(),
iter_i + epoch * epoch_size)
writer.add_scalar('local loss', txtytwth_loss.item(),
iter_i + epoch * epoch_size)
t1 = time.time()
print(
'[Epoch %d/%d][Iter %d/%d][lr %.6f]'
'[Loss: obj %.2f || cls %.2f || bbox %.2f || total %.2f || size %d || time: %.2f]'
% (epoch + 1, max_epoch, iter_i, epoch_size, tmp_lr,
conf_loss.item(), cls_loss.item(), txtytwth_loss.item(),
total_loss.item(), input_size[0], t1 - t0),
flush=True)
t0 = time.time()
# multi-scale trick
if iter_i % 10 == 0 and iter_i > 0 and args.multi_scale:
if epoch >= max_epoch - 10:
size = 608
else:
size = random.randint(10, 19) * 32
input_size = [size, size]
model.set_grid(input_size)
# change input dim
# But this operation will report bugs when we use more workers in data loader, so I have to use 0 workers.
# I don't know how to make it suit more workers, and I'm trying to solve this question.
dataloader.dataset.reset_transform(
SSDAugmentation(input_size,
mean=(0.406, 0.456, 0.485),
std=(0.225, 0.224, 0.229)))
# COCO evaluation
if (epoch + 1) % args.eval_epoch == 0:
model.trainable = False
model.set_grid(cfg['min_dim'])
# evaluate
ap50_95, ap50 = evaluator.evaluate(model)
print('ap50 : ', ap50)
print('ap50_95 : ', ap50_95)
# convert to training mode.
model.trainable = True
model.set_grid(input_size)
model.train()
if args.tfboard:
writer.add_scalar('val/COCOAP50', ap50, epoch + 1)
writer.add_scalar('val/COCOAP50_95', ap50_95, epoch + 1)
if (epoch + 1) % 10 == 0:
print('Saving state, epoch:', epoch + 1)
torch.save(
model.state_dict(),
os.path.join(path_to_save,
args.version + '_' + repr(epoch + 1) + '.pth'))
def main():
"""
YOLOv3 trainer. See README for details.
"""
args = parse_args()
print("Setting Arguments.. : ", args)
if torch.cuda.is_available() and args.gpu >= 0:
device = torch.device('cuda:{}'.format(args.gpu))
else:
device = torch.device('cpu')
os.makedirs(args.checkpoint_dir, exist_ok=True)
# Use `file_name` key to load images when train YOLO with VisDrone data
print("------------------------------------")
print(" use {} dataset for training. ".format(args.data))
print("------------------------------------")
assert args.data in ['coco', 'drone']
if args.data == 'coco':
data_dir = './COCO'
cfg_path = 'config/yolov3_default.cfg'
use_filename_key = False
if args.data == 'drone':
data_dir = './VisDrone'
cfg_path = 'config/yolov3_visdrone_default.cfg'
use_filename_key = True
# Parse config settings
with open(cfg_path, 'r') as f:
cfg = yaml.load(f)
print("successfully loaded config file: ", cfg)
lr = cfg['TRAIN']['LR']
momentum = cfg['TRAIN']['MOMENTUM']
decay = cfg['TRAIN']['DECAY']
burn_in = cfg['TRAIN']['BURN_IN']
iter_size = cfg['TRAIN']['MAXITER']
steps = eval(cfg['TRAIN']['STEPS'])
batch_size = cfg['TRAIN']['BATCHSIZE']
subdivision = cfg['TRAIN']['SUBDIVISION']
ignore_thre = cfg['TRAIN']['IGNORETHRE']
random_resize = cfg['TRAIN']['RANDRESIZE']
print('effective_batch_size = batch_size * iter_size = %d * %d' %
(batch_size, subdivision))
# Load data
print("loading dataset...")
imgsize = cfg['TRAIN']['IMGSIZE']
dataset = COCODataset(model_type=cfg['MODEL']['TYPE'],
data_dir=data_dir,
img_size=imgsize,
min_size=cfg['TRAIN']['MINSIZE'],
max_labels=cfg['TRAIN']['MAXOBJECTS'],
use_filename_key=use_filename_key,
debug=args.debug)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
shuffle=True,
num_workers=args.n_cpu)
dataiterator = iter(dataloader)
evaluator = COCOAPIEvaluator(model_type=cfg['MODEL']['TYPE'],
data_dir=data_dir,
img_size=cfg['TEST']['IMGSIZE'],
confthre=cfg['TEST']['CONFTHRE'],
nmsthre=cfg['TEST']['NMSTHRE'],
min_size=cfg['TEST']['MINSIZE'],
max_labels=cfg['TEST']['MAXOBJECTS'],
use_filename_key=use_filename_key)
# Learning rate setup
base_lr = lr
# Initiate model
n_classes = len(dataset.class_ids)
model = YOLOv3(n_classes=n_classes, ignore_thre=ignore_thre)
if args.weights_path:
print("loading darknet weights....", args.weights_path)
parse_yolo_weights(model, args.weights_path)
elif args.checkpoint:
print("loading pytorch ckpt...", args.checkpoint)
model.load_state_dict(torch.load(args.checkpoint))
print("sending model to device...")
model = model.to(device)
if args.tboard:
print("using tboard")
from tensorboardX import SummaryWriter
tblogger = SummaryWriter(os.path.join('logs', args.tboard))
model.train()
# optimizer setup
# set weight decay only on conv.weight
params_dict = dict(model.named_parameters())
params = []
for key, value in params_dict.items():
if 'conv.weight' in key:
params += [{
'params': value,
'weight_decay': decay * batch_size * subdivision
}]
else:
params += [{'params': value, 'weight_decay': 0.0}]
optimizer = optim.SGD(params,
lr=base_lr,
momentum=momentum,
dampening=0,
weight_decay=decay * batch_size * subdivision)
tmp_lr = base_lr
def set_lr(tmp_lr):
for param_group in optimizer.param_groups:
param_group['lr'] = tmp_lr / batch_size / subdivision
# start training loop
eval_interval = cfg['TRAIN']['ITER_EVAL']
checkpoint_interval = cfg['TRAIN']['ITER_CKPT']
# random resizing..
imgsize_max = imgsize
imgsize_res = 32
r_max = int(math.floor(imgsize_max / imgsize_res))
r_min = int(math.ceil(imgsize_max / imgsize_res / 2.0))
for iter_i in range(iter_size):
# COCO evaluation
if iter_i % eval_interval == 0 and iter_i > 0:
ap50_95, ap50 = evaluator.evaluate(model)
model.train()
if args.tboard:
tblogger.add_scalar('val/AP50', ap50, iter_i)
tblogger.add_scalar('val/AP50_95', ap50_95, iter_i)
# learning rate scheduling
if iter_i < burn_in:
tmp_lr = base_lr * pow(iter_i / burn_in, 4)
set_lr(tmp_lr)
elif iter_i == burn_in:
tmp_lr = base_lr
set_lr(tmp_lr)
elif iter_i in steps:
tmp_lr = tmp_lr * 0.1
set_lr(tmp_lr)
# subdivision loop
optimizer.zero_grad()
for inner_iter_i in range(subdivision):
try:
imgs, targets, _, _ = next(dataiterator) # load a batch
except StopIteration:
dataiterator = iter(dataloader)
imgs, targets, _, _ = next(dataiterator) # load a batch
imgs = Variable(imgs.to(device, dtype=torch.float32))
targets = Variable(targets.to(device, dtype=torch.float32),
requires_grad=False)
loss = model(imgs, targets)
loss.backward()
optimizer.step()
if iter_i % 10 == 0:
# logging
print(
'[Iter %d/%d] [lr %f] '
'[Losses: xy %f, wh %f, conf %f, cls %f, total %f, imgsize %d]'
% (iter_i, iter_size, tmp_lr, model.loss_dict['xy'],
model.loss_dict['wh'], model.loss_dict['conf'],
model.loss_dict['cls'], model.loss_dict['l2'], imgsize),
flush=True)
if args.tboard:
tblogger.add_scalar('train/total_loss', model.loss_dict['l2'],
iter_i)
# random resizing
if random_resize:
imgsize = random.randint(r_min, r_max) * imgsize_res
dataset.img_shape = (imgsize, imgsize)
dataset.img_size = imgsize
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
shuffle=True,
num_workers=args.n_cpu)
dataiterator = iter(dataiterator)
# save checkpoint
if iter_i > 0 and (iter_i % checkpoint_interval == 0):
torch.save(
model.state_dict(),
os.path.join(args.checkpoint_dir,
"snapshot_{}.ckpt".format(iter_i)))
if args.tboard:
tblogger.close()
def main():
"""
YOLOv3 trainer. See README for details.
"""
args = parse_args()
print("Setting Arguments.. : ", args)
cuda = torch.cuda.is_available() and args.use_cuda
os.makedirs(args.checkpoint_dir, exist_ok=True)
# Parse config settings
with open(args.cfg, 'r') as f:
cfg = yaml.load(f)
print("successfully loaded config file: ", cfg)
momentum = cfg['TRAIN']['MOMENTUM']
decay = cfg['TRAIN']['DECAY']
burn_in = cfg['TRAIN']['BURN_IN']
iter_size = cfg['TRAIN']['MAXITER']
steps = eval(cfg['TRAIN']['STEPS'])
batch_size = cfg['TRAIN']['BATCHSIZE']
subdivision = cfg['TRAIN']['SUBDIVISION']
ignore_thre = cfg['TRAIN']['IGNORETHRE']
random_resize = cfg['AUGMENTATION']['RANDRESIZE']
base_lr = cfg['TRAIN']['LR'] / batch_size / subdivision
datatype = cfg['TRAIN']['DATATYPE']
print('effective_batch_size = batch_size * iter_size = %d * %d' %
(batch_size, subdivision))
# Learning rate setup
def burnin_schedule(i):
if i < burn_in:
factor = pow(i / burn_in, 4)
elif i < steps[0]:
factor = 1.0
elif i < steps[1]:
factor = 0.1
else:
factor = 0.01
return factor
# Initiate model
model = YOLOv3(cfg['MODEL'], ignore_thre=ignore_thre)
if args.weights_path:
print("loading darknet weights....", args.weights_path)
parse_yolo_weights(model, args.weights_path)
elif args.checkpoint:
print("loading pytorch ckpt...", args.checkpoint)
state = torch.load(args.checkpoint)
if 'model_state_dict' in state.keys():
model.load_state_dict(state['model_state_dict'])
else:
model.load_state_dict(state)
if cuda:
print("using cuda")
model = model.cuda()
if args.tfboard:
print("using tfboard")
from tensorboardX import SummaryWriter
tblogger = SummaryWriter(args.tfboard)
model.train()
coco_class_names, coco_class_ids, coco_class_colors = get_coco_label_names()
imgsize = cfg['TRAIN']['IMGSIZE']
if datatype=='voc':
dataset = VOCDataset(model_type=cfg['MODEL']['TYPE'],
data_dir='../../VOCdevkit/VOC2007',
img_size=imgsize,
augmentation=cfg['AUGMENTATION'],
debug=args.debug)
print('load voc dataset successfully')
else:
dataset = COCODataset(model_type=cfg['MODEL']['TYPE'],
data_dir='COCO/',
img_size=imgsize,
augmentation=cfg['AUGMENTATION'],
debug=args.debug)
print('load COCO dataset successfully')
evaluator = COCOAPIEvaluator(model_type=cfg['MODEL']['TYPE'],
data_dir='COCO/',
img_size=cfg['TEST']['IMGSIZE'],
confthre=cfg['TEST']['CONFTHRE'],
nmsthre=cfg['TEST']['NMSTHRE'])
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=args.n_cpu)
dataiterator = iter(dataloader)
dtype = torch.cuda.FloatTensor if cuda else torch.FloatTensor
# optimizer setup
# set weight decay only on conv.weight
params_dict = dict(model.named_parameters())
params = []
for key, value in params_dict.items():
if 'conv.weight' in key:
params += [{'params': value, 'weight_decay': decay
* batch_size * subdivision}]
else:
params += [{'params': value, 'weight_decay': 0.0}]
optimizer = optim.SGD(params, lr=base_lr, momentum=momentum,
dampening=0, weight_decay=decay * batch_size * subdivision)
iter_state = 0
if args.checkpoint:
if 'optimizer_state_dict' in state.keys():
optimizer.load_state_dict(state['optimizer_state_dict'])
iter_state = state['iter'] + 1
#学习率控制 Sets the learning rate of each parameter group to the initial lr times a given function. When last_epoch=-1, sets initial lr as lr.
scheduler = optim.lr_scheduler.LambdaLR(optimizer, burnin_schedule)
# result=evals(model)
# print(result)
# start training loop
# print('args.eval_interval',args.eval_interval)
for iter_i in range(iter_state, iter_size + 1):
if iter_i % (args.eval_interval*2) == 0 and iter_i > 0:
if datatype=='voc':
result=evals(model)
print(result)
else:
ap50_95, ap50 = evaluator.evaluate(model)
print(ap50_95, ap50)
model.train()
if args.tfboard:
tblogger.add_scalar('val/COCOAP50', ap50, iter_i)
tblogger.add_scalar('val/COCOAP50_95', ap50_95, iter_i)
if iter_i % (40000) == 0 and iter_i > 0:
draw(model,datatype,imgsize)
# subdivision loop
optimizer.zero_grad()
for inner_iter_i in range(subdivision):
try:
imgs, targets, info_img, id_ = next(dataiterator) # load a batch
except StopIteration:
dataiterator = iter(dataloader)
imgs, targets, info_img, id_ = next(dataiterator) # load a batch
imgs = Variable(imgs.type(dtype))
targets = Variable(targets.type(dtype), requires_grad=False)
loss = model(imgs, targets)
loss.backward()
optimizer.step()
scheduler.step()
if iter_i % 10 == 0:
# logging
current_lr = scheduler.get_lr()[0] * batch_size * subdivision
print('[Iter %d/%d] [lr %f] '
'[Losses: xy %f, wh %f, conf %f, cls %f, total %f, imgsize %d]'
% (iter_i, iter_size, current_lr,
model.loss_dict['xy'], model.loss_dict['wh'],
model.loss_dict['conf'], model.loss_dict['cls'],
model.loss_dict['l2'], imgsize))
if args.tfboard:
tblogger.add_scalar('train/total_loss',
model.loss_dict['l2'], iter_i)
# random resizing
# 变输入大小,利用了yolov3网络的全卷积，使得模型不受图像大小的改变而影响参数。
if random_resize:
imgsize = (random.randint(0, 9) % 10 + 10) * 32
dataset.img_shape = (imgsize, imgsize)
dataset.img_size = imgsize
dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True, num_workers=args.n_cpu)
dataiterator = iter(dataloader)
if iter_i % 100 == 0:
model.eval()
if datatype=='voc':
img_file = os.path.join('../../VOCdevkit/VOC2007', 'JPEGImages', id_[0] + '.jpg')
else:
img_file = os.path.join('COCO', 'train2017',
'{:012}'.format(id_) + '.jpg')
img = cv2.imread(img_file)
img_raw = img.copy()[:, :, ::-1].transpose((2, 0, 1))
# print(img_raw.shape)
# print(img_raw)
# print(imgs)
img, info_img_t = preprocess(img, imgsize, jitter=0) # info = (h, w, nh, nw, dx, dy)
img = np.transpose(img / 255., (2, 0, 1))
img = torch.from_numpy(img).float().unsqueeze(0)
img = Variable(img.type(torch.cuda.FloatTensor))
outputs = model(img)
#outputs.shape : torch.Size([1, 12348, 85])
outputs = postprocess(outputs, 80, 0.5, 0.5)
# imgs.shape : torch.Size([1, 3, 608, 608])
# outputs[0].shape :torch.Size([3, 7])
# targets.shape :torch.Size([1, 50, 5])
# print(outputs)
if outputs[0] is not None:
bboxes = list()
classes = list()
colors = list()
# print(info_img_t)
info_img=tuple(info_img)
# print(info_img)
for x1, y1, x2, y2, conf, cls_conf, cls_pred in outputs[0]:
cls_id = coco_class_ids[int(cls_pred)]
# print(int(x1), int(y1), int(x2), int(y2), float(conf), int(cls_pred))
# print('\t+ Label: %s, Conf: %.5f' %
# (coco_class_names[cls_id], cls_conf.item()))
# print([y1, x1, y2, x2])
box = yolobox2label([y1, x1, y2, x2], info_img_t)
bboxes.append(box)
classes.append(cls_id)
colors.append(coco_class_colors[int(cls_pred)])
vis_bbox(
img_raw, bboxes, label=classes, label_names=coco_class_names,
instance_colors=colors, linewidth=2)
plt.savefig('output/'+str(iter_i)+'.jpg')
model.train()
# save checkpoint
if iter_i > 0 and (iter_i % args.checkpoint_interval == 0):
torch.save({'iter': iter_i,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
},
os.path.join(args.checkpoint_dir, "snapshot" + str(iter_i) + ".ckpt"))
if args.tfboard:
tblogger.close()
