cfg = voc_refinedet[args.input_size]
target_size = cfg['min_dim']
num_classes = cfg['num_classes']
objectness_threshold = 0.01
args.nms_threshold = 0.49 # nms
# args.nms_threshold = 0.45 # softnms
args.confidence_threshold = 0.01
args.top_k = 1000
args.keep_top_k = 500
args.vis_thres = 0.3
# args.multi_scale_test = True
set_type = 'test'
# load data
dataset = VOCDetection(args.voc_root, [('2007', set_type)],
dataset_name='VOC0712')
# load net
torch.set_grad_enabled(False)
load_to_cpu = not args.cuda
cudnn.benchmark = True
device = torch.device('cuda' if args.cuda else 'cpu')
detect = Detect_RefineDet(num_classes,
int(args.input_size),
0,
objectness_threshold,
confidence_threshold=args.confidence_threshold,
nms_threshold=args.nms_threshold,
top_k=args.top_k,
keep_top_k=args.keep_top_k)
net = build_refinedet('test', int(args.input_size), num_classes,
backbone_dict)
}
p = 0.6
img_dim = (300, 512)[args.size == '512']
bgr_means = (104, 117, 123)
num_classes = (21, 81)[args.dataset == 'COCO']
batch_size = args.batch_size
weight_decay = 0.0005
gamma = 0.1
momentum = 0.9
dataset_name = args.dataset
if dataset_name[0] == "V":
cfg = (VOC_300, VOC_512)[args.size == '512']
train_dataset = VOCDetection(VOCroot, datasets_dict[dataset_name],
SSDAugmentation(img_dim, bgr_means),
AnnotationTransform(), dataset_name)
# train_dataset = VOCDetection(VOCroot, datasets_dict[dataset_name], preproc(img_dim, bgr_means, p), AnnotationTransform())
test_dataset = VOCDetection(VOCroot, datasets_dict["VOC2007"], None,
AnnotationTransform(), dataset_name)
elif dataset_name[0] == "C":
train_dataset = COCODetection(COCOroot, datasets_dict[dataset_name],
SSDAugmentation(img_dim, bgr_means),
COCOAnnotationTransform(), dataset_name)
test_dataset = COCODetection(COCOroot, datasets_dict["COCOval"], None,
COCOAnnotationTransform(), dataset_name)
cfg = (COCO_300, COCO_512)[args.size == '512']
else:
print('Unkown dataset!')
if args.version == "ssd_vgg":
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == 'module.':
name = k[7:] # remove `module.`
else:
name = k
new_state_dict[name] = v
net.load_state_dict(new_state_dict)
net.eval()
print('Finished loading model!')
print(net)
# load data
if args.dataset == 'VOC':
testset = VOCDetection(
VOCroot, [('2007', 'test')], None, AnnotationTransform())
elif args.dataset == 'COCO':
testset = COCODetection(
COCOroot, [('2014', 'minival')], None)
#COCOroot, [('2015', 'test-dev')], None)
else:
print('Only VOC and COCO dataset are supported now!')
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
else:
net = net.cpu()
# evaluation
#top_k = (300, 200)[args.dataset == 'COCO']
top_k = 200
detector = Detect(num_classes,0,cfg)
def evaluate_detections(box_list, output_dir, dataset):
write_voc_results_file(box_list, dataset)
do_python_eval(output_dir)
if __name__ == '__main__':
# load net
num_classes = len(labelmap) + 1 # +1 for background
net = build_ssd('test', 300, num_classes) # initialize SSD
net.load_state_dict(torch.load(args.trained_model))
net.eval()
print('Finished loading model!')
# load data
dataset = VOCDetection(args.voc_root, [('2007', set_type)],
BaseTransform(300, dataset_mean),
VOCAnnotationTransform())
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
# evaluation
test_net(args.save_folder,
net,
args.cuda,
dataset,
BaseTransform(net.size, dataset_mean),
args.top_k,
300,
thresh=args.confidence_threshold)
def train():
net.train()
# loss counters
loc_loss = 0 # epoch
conf_loss = 0
epoch = 0
print('Loading Dataset...')
if args.datasets == 'voc':
train_sets = [('2007', 'trainval'), ('2012', 'trainval')]
dataset = VOCDetection(args.voc_root, train_sets, SSDAugmentation(
ssd_dim, means), AnnotationTransform())
elif args.datasets == 'coco':
dataset = make_dataset('coco',args.coco_dataroot, args.coco_annfile)
print(len(dataset))
epoch_size = len(dataset) // args.batch_size
#print('Training SSD on', dataset.name)
step_index = 0
if args.visdom:
# initialize visdom loss plot
lot = viz.line(
X=torch.zeros((1,)).cpu(),
Y=torch.zeros((1, 3)).cpu(),
opts=dict(
xlabel='Iteration',
ylabel='Loss',
title='Current SSD Training Loss',
legend=['Loc Loss', 'Conf Loss', 'Loss']
)
)
epoch_lot = viz.line(
X=torch.zeros((1,)).cpu(),
Y=torch.zeros((1, 3)).cpu(),
opts=dict(
xlabel='Epoch',
ylabel='Loss',
title='Epoch SSD Training Loss',
legend=['Loc Loss', 'Conf Loss', 'Loss']
)
)
batch_iterator = None
data_loader = data.DataLoader(dataset, batch_size, num_workers=args.num_workers,
shuffle=True, pin_memory=True, collate_fn = detection_collate)
scheduler = ReduceLROnPlateau(optimizer, 'min', min_lr = 1e-10, verbose = True)
for iteration in range(args.start_iter, max_iter):
if (not batch_iterator) or (iteration % epoch_size == 0):
# create batch iterator
batch_iterator = iter(data_loader)
if iteration in stepvalues:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
if args.visdom:
viz.line(
X=torch.ones((1, 3)).cpu() * epoch,
Y=torch.Tensor([loc_loss, conf_loss,
loc_loss + conf_loss]).unsqueeze(0).cpu() / epoch_size,
win=epoch_lot,
update='append'
)
# reset epoch loss counters
loc_loss = 0
conf_loss = 0
epoch += 1
# load train data
images, targets = next(batch_iterator)
if args.cuda:
images = Variable(images.cuda())
targets = [Variable(anno.cuda(), volatile=True) for anno in targets]
else:
images = Variable(images)
targets = [Variable(anno, volatile=True) for anno in targets]
# forward
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, targets)
loss = 0.01*loss_l + loss_c
loss.backward()
optimizer.step()
t1 = time.time()
loc_loss += loss_l.data[0]
conf_loss += loss_c.data[0]
#scheduler.step(loss.data[0])
if iteration % 10 == 0:
print('Timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || Loss_loc: %.4f Loss_conf %.4f ||' % (loss_l.data[0], loss_c.data[0]))
if args.visdom and args.send_images_to_visdom:
random_batch_index = np.random.randint(images.size(0))
viz.image(images.data[random_batch_index].cpu().numpy())
if args.visdom:
viz.line(
X=torch.ones((1, 3)).cpu() * iteration,
Y=torch.Tensor([loss_l.data[0], loss_c.data[0],
loss_l.data[0] + loss_c.data[0]]).unsqueeze(0).cpu(),
win=lot,
update='append'
)
# hacky fencepost solution for 0th epoch plot
if iteration == 0:
viz.line(
X=torch.zeros((1, 3)).cpu(),
Y=torch.Tensor([loc_loss, conf_loss,
loc_loss + conf_loss]).unsqueeze(0).cpu(),
win=epoch_lot,
update=True
)
if iteration % 10000 == 0:
print('Saving state, iter:', iteration)
save_dir = os.path.join(args.save_folder, args.experiments)
if not os.path.exists(save_dir):
os.mkdir(save_dir)
torch.save(ssd_net.state_dict(), os.path.join(save_dir,'ssd300_0712_' +
repr(iteration) + '.pth'))
torch.save(ssd_net.state_dict(), args.save_folder + '' + args.version + '.pth')
ymax = top_ymax[j] * img.size[1]
score = top_conf[j]
label_name = top_labels[j]
# print(img_id)
coords = [xmin, ymin, xmax - xmin, ymax - ymin, score, img_id]
with open(filename, mode='a') as f:
f.write('PREDICTION: ' + repr(j) + '\n')
f.write('label: ' + label_name + ' score: ' + str(score) +
' ' + ' || '.join(str(c) for c in coords) + '\n\n')
if __name__ == '__main__':
# load net
net = build_ssd('test', 300, 21) # initialize SSD
net.load_state_dict(torch.load(args.trained_model))
net.eval()
print('Finished loading model!')
# load data
valset = VOCDetection(VOCroot, 'val', None, AnnotationTransform())
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
# evaluation
test_net(args.save_folder,
net,
args.cuda,
valset,
base_transform(net.size, (104, 117, 123)),
args.top_k,
thresh=args.confidence_threshold)
def train(model, resume=False):
model.train()
optimizer = build_optimizer(args, model)
scheduler = build_lr_scheduler(args, optimizer)
checkpointer = DetectionCheckpointer(
model, args, optimizer=optimizer, scheduler=scheduler
)
criterion = MultiBoxLoss_combined(num_classes, overlap_threshold, True, 0, True, 3, 0.5, False)
start_iter = (
checkpointer.resume_or_load(args.basenet if args.phase == 1 else args.load_file,
resume=resume).get("iteration", -1) + 1
)
max_iter = args.max_iter
periodic_checkpointer = PeriodicCheckpointer(
checkpointer, args.checkpoint_period, max_iter=max_iter
)
writers = (
[
CommonMetricPrinter(max_iter),
TensorboardXWriter(args.save_folder),
]
)
if args.dataset == 'VOC':
dataset = VOCDetection(args, VOCroot, train_sets, preproc(
img_dim, rgb_means, p), AnnotationTransform(0 if args.setting == 'transfer' else args.split))
elif args.dataset == 'COCO':
dataset = COCODetection(COCOroot, train_sets, preproc(
img_dim, rgb_means, p))
else:
raise ValueError(f"Unknown dataset: {args.dataset}")
if args.phase == 2 and args.method == 'ours':
sampler = TrainingSampler(len(dataset))
data_loader = torch.utils.data.DataLoader(
dataset,
args.batch_size,
sampler=sampler,
num_workers=args.num_workers,
collate_fn=detection_collate,
)
# initialize the OBJ(Target) parameters
init_reweight(args, model, data_loader)
dataset.set_mixup(np.random.beta, 1.5, 1.5)
logger.info('Fine tuning on ' + str(args.shot) + '-shot task')
sampler = TrainingSampler(len(dataset))
data_loader = iter(torch.utils.data.DataLoader(
dataset,
args.batch_size,
sampler=sampler,
num_workers=args.num_workers,
collate_fn=detection_collate,
))
assert model.training, 'Model.train() must be True during training.'
logger.info("Starting training from iteration {}".format(start_iter))
# scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones, gamma=args.gamma, last_epoch=epoch - 1)
with EventStorage(start_iter) as storage:
for iteration in range(start_iter, max_iter):
iteration = iteration + 1
storage.step()
if args.phase == 2 and args.method == 'ours' and \
iteration == (args.max_iter - args.no_mixup_iter):
dataset.set_mixup(None)
data_loader = iter(torch.utils.data.DataLoader(
dataset,
args.batch_size,
sampler=sampler,
num_workers=args.num_workers,
collate_fn=detection_collate,
))
data, targets = next(data_loader)
# storage.put_image('image', vis_tensorboard(data))
output = model(data)
loss_dict = criterion(output, priors, targets)
losses = sum(loss for loss in loss_dict.values())
# assert torch.isfinite(losses).all(), loss_dict
storage.put_scalars(total_loss=losses, **loss_dict)
optimizer.zero_grad()
losses.backward()
optimizer.step()
if args.phase == 2 and args.method == 'ours':
if isinstance(model, (DistributedDataParallel, DataParallel)):
model.module.normalize()
else:
model.normalize()
storage.put_scalar("lr", optimizer.param_groups[-1]["lr"], smoothing_hint=False)
scheduler.step()
if iteration - start_iter > 5 and (iteration % 20 == 0 or iteration == max_iter):
for writer in writers:
writer.write()
periodic_checkpointer.step(iteration)
def train():
net.train()
epoch = args.start_iter
if args.dataset_name == 'COCO':
dataset = COCODetection(COCOroot,
year='trainval2014',
image_sets=train_sets,
transform=data_transform(ssd_dim, means),
phase='train')
else:
dataset = VOCDetection(
data_root,
train_sets,
data_transform(ssd_dim, means),
AnnotationTransform(dataset_name=args.dataset_name),
dataset_name=args.dataset_name,
set_file_name=set_filename)
epoch_size = len(dataset) // args.batch_size
drop_step = [s * epoch_size for s in args.step_list]
max_iter = max_epoch * epoch_size
logging.info('Loading Dataset:' + args.dataset_name + ' dataset size: ' +
str(len(dataset)))
step_index = 0
if args.visdom:
# initialize visdom loss plot
y_dim = 3
legend = [
'Loss',
'Loc Loss',
'Conf Loss',
]
if use_refine:
y_dim += 1
legend += [
'Arm Loc Loss',
]
lot = viz.line(X=torch.zeros((1, )),
Y=torch.zeros((1, y_dim)),
opts=dict(
xlabel='Iteration',
ylabel='Loss',
title=args.save_folder.split('/')[-1],
legend=legend,
))
batch_iterator = None
data_loader = data.DataLoader(dataset,
batch_size,
num_workers=args.num_workers,
shuffle=True,
collate_fn=collate_fn,
pin_memory=True)
for iteration in range(epoch * epoch_size, max_iter + 10):
if (not batch_iterator) or (iteration % epoch_size == 0):
# create batch iterator
batch_iterator = iter(data_loader)
if epoch % args.save_interval == 0:
logging.info('Saving state, epoch: ' + str(epoch))
torch.save(
ssd_net.state_dict(),
os.path.join(
args.save_folder, args.model_name + str(ssd_dim) +
'_' + args.dataset_name + '_' + repr(epoch) + '.pth'))
epoch += 1
t0 = time.time()
if iteration in drop_step:
step_index = drop_step.index(iteration) + 1
adjust_learning_rate(optimizer, args.gamma, epoch, step_index,
iteration, epoch_size)
# adjust_learning_rate(optimizer, args.gamma)
collected_data = next(batch_iterator)
with torch.no_grad():
images, targets = collected_data[:2]
images = images.to(device)
targets = [anno.to(device) for anno in targets]
# forward
loss = torch.tensor(0., requires_grad=True).to(device)
out = net(images)
# backward
optimizer.zero_grad()
if use_refine:
loss_arm_l = arm_criterion(out[0], priors, targets)
loss_l, loss_c = criterion(out[2:],
priors,
targets,
arm_data=out[:2])
loss += args.loss_coe[0] * loss_arm_l
else:
loss_l, loss_c = criterion(out, priors, targets)
loss += args.loss_coe[0] * loss_l + args.loss_coe[1] * loss_c
loss.backward()
optimizer.step()
t1 = time.time()
if iteration % 10 == 0:
if use_refine:
logging.info(
'Epoch:' + repr(epoch) + ', epochiter: ' +
repr(iteration % epoch_size) + '/' + repr(epoch_size) +
', total_iter ' + repr(iteration) +
' || loss: %.4f, Loss_l: %.4f, loss_c: %.4f, loss_arm_l: %.4f, lr: %.5f || Timer: %.4f sec.'
% (loss, loss_l, loss_c, loss_arm_l,
optimizer.param_groups[0]['lr'], t1 - t0))
else:
logging.info(
'Epoch:' + repr(epoch) + ', epochiter: ' +
repr(iteration % epoch_size) + '/' + repr(epoch_size) +
', total_iter ' + repr(iteration) +
' || loss: %.4f, Loss_l: %.4f, loss_c: %.4f, lr: %.5f || Timer: %.4f sec.'
% (loss, loss_l, loss_c, optimizer.param_groups[0]['lr'],
t1 - t0))
if args.visdom:
y_dis = [
loss.cpu(), args.loss_coe[0] * loss_l.cpu(),
args.loss_coe[1] * loss_c.cpu()
]
if iteration == 1000:
# initialize visdom loss plot
lot = viz.line(X=torch.zeros((1, )),
Y=torch.zeros((1, y_dim)),
opts=dict(
xlabel='Iteration',
ylabel='Loss',
title=args.save_folder.split('/')[-1],
legend=legend,
))
if use_refine:
y_dis += [
args.loss_coe[0] * loss_arm_l.cpu(),
]
# update = 'append' if iteration
viz.line(X=torch.ones((1, y_dim)) * iteration,
Y=torch.FloatTensor(y_dis).unsqueeze(0),
win=lot,
update='append',
opts=dict(
xlabel='Iteration',
ylabel='Loss',
title=args.save_folder.split('/')[-1],
legend=legend,
))
torch.save(
ssd_net.state_dict(),
os.path.join(
args.save_folder, args.model_name + str(ssd_dim) + '_' +
args.dataset_name + '_' + repr(iteration) + '.pth'))
print('Complet Training. Saving state, iter:', iteration)
def train():
net.train()
# loss counters
loc_loss = 0 # epoch
conf_loss = 0
epoch = 0
print('Loading Dataset...')
dataset = VOCDetection(VOCroot, train_sets, preproc(img_dim, rgb_means, p),
AnnotationTransform())
epoch_size = len(dataset) // args.batch_size
max_iter = args.max_epoch * epoch_size
stepvalues = (25 * epoch_size, 35 * epoch_size)
step_index = 0
start_iter = 0
lr = args.lr
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
if (epoch > 20 and epoch % 2 == 0):
torch.save(
net.state_dict(), args.save_folder + 'epoches_' +
repr(epoch).zfill(3) + '.pth')
# create batch iterator
batch_iterator = iter(
data.DataLoader(dataset,
batch_size,
shuffle=True,
num_workers=8,
collate_fn=detection_collate))
loc_loss = 0
conf_loss = 0
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(optimizer, 0.2, epoch, step_index, iteration,
epoch_size)
images, targets = next(batch_iterator)
images = Variable(images.cuda())
targets = [Variable(anno.cuda()) for anno in targets]
# forward
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, priors, targets)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
t1 = time.time()
loc_loss += loss_l.item()
conf_loss += loss_c.item()
load_t1 = time.time()
# visualization
visualize_total_loss(writer, loss.item(), iteration)
visualize_loc_loss(writer, loss_l.item(), iteration)
visualize_conf_loss(writer, loss_c.item(), iteration)
if iteration % 10 == 0:
print('Epoch:' + repr(epoch) + ' || epochiter: ' +
repr(iteration % epoch_size) + '/' + repr(epoch_size) +
'|| Totel iter ' + repr(iteration) +
' || L: %.4f C: %.4f||' % (loss_l.item(), loss_c.item()) +
'Batch time: %.4f sec. ||' % (load_t1 - load_t0) +
'LR: %.8f' % (lr))
torch.save(net.state_dict(),
args.save_folder + 'epoches_' + repr(epoch).zfill(3) + '.pth')
# load net
if args.dataset == 'VOC':
cfg = voc
else:
cfg = custom
# +1 for background
net = build_ssd(phase='test', size=cfg['min_dim'],
num_classes=cfg['num_classes'], confidence_threshold=args.confidence_threshold)
net.load_state_dict(torch.load(args.trained_model))
net = net.to(device)
net.eval()
print('Finished loading model!')
# Load data (TODO: add COCO)
if args.dataset == 'VOC':
dataset = VOCDetection(args.dataset_root, [(set_type)],
BaseTransform(cfg['min_dim'], MEANS),
VOCAnnotationTransform())
else:
dataset = CustomDetection(root=args.dataset_root,
image_set=[(set_type)],
transform=BaseTransform(cfg['min_dim'], MEANS),
target_transform=CustomAnnotationTransform(train=False))
# net = net.to(device)
# Evaluation
test_net(args.save_folder, net, args.cuda, dataset,
BaseTransform(net.size, MEANS), args.top_k, cfg['min_dim'],
thresh=args.confidence_threshold)
def train():
net.train()
# loss counters
loc_loss = 0 # epoch
conf_loss = 0
epoch = 0
print('Loading Dataset...')
dataset = VOCDetection(VOCroot, train_sets, BaseTransform(
ssd_dim, rgb_means), AnnotationTransform())
epoch_size = len(dataset) // args.batch_size
print('Training SSD on', dataset.name)
step_index = 0
if args.visdom:
# initialize visdom loss plot
lot = viz.line(
X=torch.zeros((1,)).cpu(),
Y=torch.zeros((1, 3)).cpu(),
opts=dict(
xlabel='Iteration',
ylabel='Loss',
title='Current SSD Training Loss',
legend=['Loc Loss', 'Conf Loss', 'Loss']
)
)
epoch_lot = viz.line(
X=torch.zeros((1,)).cpu(),
Y=torch.zeros((1, 3)).cpu(),
opts=dict(
xlabel='Epoch',
ylabel='Loss',
title='Epoch SSD Training Loss',
legend=['Loc Loss', 'Conf Loss', 'Loss']
)
)
for iteration in range(max_iter):
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(data.DataLoader(dataset, batch_size,
shuffle=True, collate_fn=detection_collate))
if iteration in stepvalues:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
if args.visdom:
viz.line(
X=torch.ones((1, 3)).cpu() * epoch,
Y=torch.Tensor([loc_loss, conf_loss,
loc_loss + conf_loss]).unsqueeze(0).cpu() / epoch_size,
win=epoch_lot,
update='append'
)
# reset epoch loss counters
loc_loss = 0
conf_loss = 0
epoch += 1
# load train data
images, targets = next(batch_iterator)
# print(images)
# print(targets)
if args.cuda:
images = Variable(images.cuda())
targets = [Variable(anno.cuda()) for anno in targets]
else:
images = Variable(images)
targets = [Variable(anno) for anno in targets]
# forward
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, targets)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
t1 = time.time()
loc_loss += loss_l.data[0]
conf_loss += loss_c.data[0]
if iteration % 10 == 0:
print('Timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.data[0]), end=' ')
if args.visdom:
viz.line(
X=torch.ones((1, 3)).cpu() * iteration,
Y=torch.Tensor([loss_l.data[0], loss_c.data[0],
loss_l.data[0] + loss_c.data[0]]).unsqueeze(0).cpu(),
win=lot,
update='append'
)
# hacky fencepost solution for 0th epoch plot
if iteration == 0:
viz.line(
X=torch.zeros((1, 3)).cpu(),
Y=torch.Tensor([loc_loss, conf_loss,
loc_loss + conf_loss]).unsqueeze(0).cpu(),
win=epoch_lot,
update=True
)
if iteration % 5000 == 0:
torch.save(net.state_dict(), 'weights/ssd300_0712_iter_' +
repr(iteration) + '.pth')
torch.save(net.state_dict(), args.save_folder + '' + args.version + '.pth')
if __name__ == "__main__":
# init xmls
init_xmls(args.image_root, args.annotation_root, TEMP_PATH)
# load net
num_classes = len(labelmap) + 1 # +1 for background
net = build_ssd("test", 300, num_classes) # initialize SSD
net.load_state_dict(torch.load(args.trained_model))
net.eval()
print("Finished loading model!")
# load data
dataset = VOCDetection(
args.image_root,
args.annotation_root,
[("", set_type)],
BaseTransform(300, dataset_mean),
VOCAnnotationTransform(),
)
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
# evaluation
test_net(
args.save_folder,
net,
args.cuda,
dataset,
BaseTransform(net.size, dataset_mean),
args.top_k,
300,
from loss import MultiBoxLoss
from data import VOCDetection, detection_collate
from params import voc_config, args, MEANS, VOC_ROOT
import time
import torch.optim as optim
import torch.backends.cudnn as cudnn
import torch.utils.data as data
import torch
import numpy as np
from net import build_ssd
from matplotlib import pyplot as plt
dataset = VOCDetection(root=VOC_ROOT, transform=SSDAugmentation(300, MEANS))
# ネットワークの定義
ssd_net = build_ssd('train', voc_config['num_classes'])
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# ネットワークをGPUに転送
net = ssd_net.to(device)
# 学習の再開時はargs['resume']のパラメータをロード
if args['resume']:
print('Resuming training, loading {}...'.format(args['resume']))
ssd_net.load_weights(args['save_folder'] + args['resume'])
# 新規モデル学習時は追加ネットワークの畳み込み、位置推定、クラス分類の畳み込みパラメータを初期化する
net.init_weights(args['resume'])
def test():
print('Loading Dataset...')
if args.dataset == 'VOC':
testset = VOCDetection(VOCroot, [('2007', 'test')], None,
AnnotationTransform())
train_dataset = VOCDetection(VOCroot, train_sets,
preproc(img_dim, rgb_means, rgb_std, p),
AnnotationTransform())
# elif args.dataset == 'COCO':
# testset = COCODetection(
# COCOroot, [('2017', 'val')], None)
# train_dataset = COCODetection(COCOroot, train_sets, preproc(
# img_dim, rgb_means, rgb_std, p))
else:
print('Only VOC and COCO are supported now!')
exit()
net.eval()
# loss counters
epoch = 0
if args.resume_net:
epoch = 0 + args.resume_epoch
epoch_size = len(train_dataset) // args.batch_size
max_iter = args.max_epoch * epoch_size
stepvalues_VOC = (150 * epoch_size, 200 * epoch_size, 250 * epoch_size)
stepvalues_COCO = (90 * epoch_size, 120 * epoch_size, 140 * epoch_size)
stepvalues = (stepvalues_VOC, stepvalues_COCO)[args.dataset == 'COCO']
print('Training', args.version, 'on', train_dataset.name)
step_index = 0
start_iter = 0
log_file = open(log_file_path, 'w')
batch_iterator = None
mean_loss_c = 0
mean_loss_l = 0
for iteration in range(start_iter, max_iter + 10):
if (iteration % epoch_size == 0):
# create batch iterator
batch_iterator = iter(
data.DataLoader(train_dataset,
batch_size,
shuffle=True,
num_workers=args.num_workers,
collate_fn=detection_collate))
top_k = (300, 200)[args.dataset == 'COCO']
if args.dataset == 'VOC':
APs, mAP = test_net(test_save_dir,
net,
detector,
args.cuda,
testset,
BaseTransform(net.size, rgb_means, rgb_std,
(2, 0, 1)),
top_k,
thresh=0.01)
APs = [str(num) for num in APs]
mAP = str(mAP)
log_file.write(str(iteration) + ' APs:\n' + '\n'.join(APs))
log_file.write('mAP:\n' + mAP + '\n')
else:
test_net(test_save_dir,
net,
detector,
args.cuda,
testset,
BaseTransform(net.module.size, rgb_means, rgb_std,
(2, 0, 1)),
top_k,
thresh=0.01)
load_t0 = time.time()
if iteration in stepvalues:
step_index = stepvalues.index(iteration) + 1
# load train data
images, targets = next(batch_iterator)
if args.cuda:
images = images.cuda()
targets = [anno.cuda() for anno in targets]
# forward
out = net(images)
# backprop
load_t1 = time.time()
if iteration % 2 == 0:
print('Epoch:' + repr(epoch) + ' batch: ' +
repr(iteration % epoch_size) + '/' + repr(epoch_size) +
' loss-L: %.4f loss-C: %.4f' %
(mean_loss_l / 10, mean_loss_c / 10) + 'time: %.4f ' %
(load_t1 - load_t0))
log_file.write('Epoch:' + repr(epoch) + ' epochiter: ' +
repr(iteration % epoch_size) + '/' +
repr(epoch_size) + ' Totel iter ' +
repr(iteration) + ' Loss %.4f C: %.4f' %
(mean_loss_l / 10, mean_loss_c / 10) +
'Batch time: %.4f sec. ||' % (load_t1 - load_t0) +
'\n')
mean_loss_c = 0
mean_loss_l = 0
if args.visdom and args.send_images_to_visdom:
random_batch_index = np.random.randint(images.size(0))
viz.image(images.data[random_batch_index].cpu().numpy())
log_file.close()
devices = [0, 1]
device, device_ids = prepare_device(devices)
T_file = 'model/KD_' + repr(Feature_num) + 'F' + repr(Feature_num) + 'D/T_const/prune_list/' + repr(
Epoch) + 'epoch/lr_040812/20200404/'
if os.path.exists(Path(T_file)) is False:
os.makedirs(T_file)
priorbox = PriorBox(voc, device).to(device)
priorbox = torch.nn.DataParallel(priorbox, device_ids=device_ids)
T_model = build_ssd('train', 300, 21, device)
try:
dataset = VOCDetection(root='/home/hanchengye/data/VOCdevkit/',
image_sets=[('2007', 'trainval'), ('2012', 'trainval')],
transform=SSDAugmentation(300, MEANS))
except FileNotFoundError:
dataset = VOCDetection(root='/remote-home/source/remote_desktop/share/Dataset/VOCdevkit/',
image_sets=[('2007', 'trainval'), ('2012', 'trainval')],
transform=SSDAugmentation(300, MEANS))
data_loader_rank = data.DataLoader(dataset, batch_size=32,
num_workers=4, shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
param_T = torch.load('weights/ssd300_COCO_395000.pth')
T_model.load_state_dict(param_T)
T_model = T_model.to(device)
S_model = copy.deepcopy(T_model)
def train():
net.train()
# loss counters
loc_loss = 0 # epoch
conf_loss = 0
epoch = 0 + args.resume_epoch
print('Loading Dataset...')
if args.dataset == 'VOC':
dataset = VOCDetection(VOCroot, train_sets,
preproc(img_dim, rgb_means, p),
AnnotationTransform())
elif args.dataset == 'COCO':
dataset = COCODetection(COCOroot, train_sets,
preproc(img_dim, rgb_means, p))
else:
print('Only VOC and COCO are supported now!')
return
epoch_size = len(dataset) // args.batch_size
max_iter = args.max_epoch * epoch_size
stepvalues_VOC = (150 * epoch_size, 200 * epoch_size, 250 * epoch_size)
stepvalues_COCO = (90 * epoch_size, 120 * epoch_size, 140 * epoch_size,
215 * epoch_size)
if args.size == '512' and args.batch_size < 32:
# batchsize => 8
stepvalues_COCO = (90 * epoch_size, 120 * epoch_size, 140 * epoch_size,
215 * epoch_size)
print("!!!!!!!!!")
stepvalues = (stepvalues_VOC, stepvalues_COCO)[args.dataset == 'COCO']
print('Training', args.version, 'on', dataset.name)
step_index = 0
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
lr = args.lr
loc_loss = 0
conf_loss = 0
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(
data.DataLoader(dataset,
batch_size,
shuffle=True,
num_workers=args.num_workers,
collate_fn=detection_collate))
loc_loss = 0
conf_loss = 0
if (epoch % 10 == 0 and epoch > 0) or (epoch % 1 == 0
and epoch > 90):
torch.save(
net.state_dict(), args.save_folder + args.version + '_' +
args.dataset + '_epoches_' + repr(epoch) + '.pth')
epoch += 1
load_t0 = time.time()
step_index = 0
for stepvalue in stepvalues:
if iteration >= stepvalue:
step_index += 1
lr = adjust_learning_rate(optimizer, args.gamma, epoch, step_index,
iteration, epoch_size)
images, targets = next(batch_iterator)
# load train data
lam = np.random.beta(0.1, 0.1)
## mixup
index = np.arange(len(targets))
np.random.shuffle(index)
# print(np.sum([torch.sum(anno[:,-1] == 2) for anno in targets]))
if args.cuda:
images = Variable(images.cuda())
targets = [Variable(anno.cuda()) for anno in targets]
else:
images = Variable(images)
targets = [Variable(anno) for anno in targets]
mixed_x = lam * images + (1 - lam) * images[index, :]
y_a = targets
y_b = [targets[idx] for idx in index]
# forward
t0 = time.time()
out = net(mixed_x)
# backprop
optimizer.zero_grad()
# loss_l, loss_c = criterion(out, priors, targets)
loss_l1, loss_c1 = criterion(out, priors, y_a)
loss_l2, loss_c2 = criterion(out, priors, y_b)
loss_l = lam * loss_l1 + (1 - lam) * loss_l2
loss_c = lam * loss_c1 + (1 - lam) * loss_c2
loss = loss_l + loss_c
loss.backward()
optimizer.step()
t1 = time.time()
loc_loss += loss_l.item()
conf_loss += loss_c.item()
load_t1 = time.time()
if iteration % 10 == 0:
print('Epoch:' + repr(epoch) + ' || epochiter: ' +
repr(iteration % epoch_size) + '/' + repr(epoch_size) +
'|| Totel iter ' + repr(iteration) +
' || L: %.4f C: %.4f||' % (loss_l.item(), loss_c.item()) +
'Batch time: %.4f sec. ||' % (load_t1 - load_t0) +
'LR: %.8f' % (lr))
if iteration % epoch_size == 0:
# Summary
writer.add_scalar("conf_loss", conf_loss, iteration // epoch_size)
writer.add_scalar("loc_loss", loc_loss, iteration // epoch_size)
writer.add_scalar("learning rate", lr, iteration // epoch_size)
torch.save(
net.state_dict(), args.save_folder + 'Final_' + args.version + '_' +
args.dataset + '.pth')
def train():
net.train()
epoch = 0 + args.resume_epoch
print('Loading Dataset...')
dataset = VOCDetection(args.training_dataset, preproc(img_dim, rgb_means),
AnnotationTransform())
epoch_size = math.ceil(len(dataset) / args.batch_size)
max_iter = args.max_epoch * epoch_size
stepvalues = (200 * epoch_size, 250 * epoch_size)
step_index = 0
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(
data.DataLoader(dataset,
batch_size,
shuffle=True,
num_workers=args.num_workers,
collate_fn=detection_collate))
if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 == 0
and epoch > 200):
torch.save(
net.state_dict(), args.save_folder + 'FaceBoxes_epoch_' +
repr(epoch) + '.pth')
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(optimizer, args.gamma, epoch, step_index,
iteration, epoch_size)
# load train data
images, targets = next(batch_iterator)
if gpu_train:
images = Variable(images.cuda())
targets = [Variable(anno.cuda()) for anno in targets]
else:
images = Variable(images)
targets = [Variable(anno) for anno in targets]
# forward
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, priors, targets)
loss = cfg['loc_weight'] * loss_l + loss_c
loss.backward()
optimizer.step()
load_t1 = time.time()
print('Epoch:' + repr(epoch) + ' || epochiter: ' +
repr(iteration % epoch_size) + '/' + repr(epoch_size) +
'|| Totel iter ' + repr(iteration) + ' || L: %.4f C: %.4f||' %
(cfg['loc_weight'] * loss_l.item(), loss_c.item()) +
'Batch time: %.4f sec. ||' % (load_t1 - load_t0) + 'LR: %.8f' %
(lr))
torch.save(net.state_dict(), args.save_folder + 'Final_FaceBoxes.pth')
logging.info("==> save eval result txt of each class to the directory of: ",
eval_result_txt)
if not os.path.exists(eval_result_txt):
os.makedirs(eval_result_txt)
if not os.path.exists(cachedir):
os.makedirs(cachedir)
logging.info("==> save annotation cache file to the directory of: ", cachedir)
output_dir = get_output_dir(save_folder, "PR_curves")
logging.info("==> save detections.pkl to the directory of: ", output_dir)
comp_id = "comp4_"
dataset = VOCDetection(VOCroot, [(set_type)],
BaseTransform(input_size, dataset_mean),
AnnotationTransform())
if use_cuda:
net = net.cuda()
class Timer(object):
"""A simple timer."""
def __init__(self):
self.total_time = 0.
self.calls = 0
self.start_time = 0.
self.diff = 0.
self.average_time = 0.
def tic(self):
detector = Detect(num_classes, 0, cfg)
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
criterion = MultiBoxLoss(num_classes, 0.5, True, 0, True, 3, 0.5, False)
priorbox = PriorBox(cfg)
with torch.no_grad():
priors = priorbox.forward()
# -------------------------------------------------------------------------------------------------------------------- #
# -------------------------------------------------------------------------------------------------------------------- #
# Load dataset
print('Loading Dataset...')
if args.dataset == 'VOC':
testset = VOCDetection(VOCroot, [('2007', 'test')], None,
AnnotationTransform())
train_dataset = VOCDetection(VOCroot, train_sets,
preproc(img_dim, rgb_means, rgb_std, p),
AnnotationTransform())
elif args.dataset == 'COCO':
testset = COCODetection(COCOroot, [('2017', 'val')], None)
train_dataset = COCODetection(COCOroot, train_sets,
preproc(img_dim, rgb_means, rgb_std, p))
else:
print('Only VOC and COCO are supported now!')
exit()
# -------------------------------------------------------------------------------------------------------------------- #
# -------------------------------------------------------------------------------------------------------------------- #
from utils.tuto4_logger import Logger_tf
logger = Logger_tf(save_folder)
def train():
args = parse_args()
if not os.path.exists(args.save_folder):
os.mkdir(args.save_folder)
if args.dataset == 'VOC':
train_sets = [('2007', 'trainval'), ('2012', 'trainval')]
cfg = (VOC_300, VOC_512)[args.size == '512']
else:
train_sets = [('2017', 'train'), ('2014', 'val')]
cfg = (COCO_300, COCO_512)[args.size == '512']
if args.version == 'RFB_vgg':
from models.RFB_Net_vgg import build_rfb_vgg_net
elif args.version == 'RFB_E_vgg':
from models.RFB_Net_E_vgg import build_net
elif args.version == 'RFB_mobile':
from models.RFB_Net_mobile import build_rfb_mobilenet
cfg = COCO_mobile_300
else:
print('Unkown version!')
img_dim = (300, 512)[args.size == '512']
rgb_means = ((104, 117, 123), (103.94, 116.78, 123.68))[args.version == 'RFB_mobile']
p = (0.6, 0.2)[args.version == 'RFB_mobile']
num_classes = (21, 81)[args.dataset == 'COCO']
batch_size = args.batch_size
print('Loading Dataset...')
if args.dataset == 'VOC':
dataset = VOCDetection(VOCroot, train_sets, preproc(img_dim, rgb_means, p), AnnotationTransform())
elif args.dataset == 'COCO':
dataset = COCODetection(COCOroot, train_sets, preproc(img_dim, rgb_means, p))
else:
print('Only VOC and COCO are supported now!')
return
net = build_rfb_mobilenet('train', img_dim, num_classes)
if args.resume_net:
# load resume network
print('Loading resume network...')
state_dict = torch.load(args.resume_net)
# create new OrderedDict that does not contain `module.`
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == 'module.':
name = k[7:] # remove `module.`
else:
name = k
new_state_dict[name] = v
net.load_state_dict(new_state_dict)
print('Resume model load done.')
else:
base_weights = torch.load(args.basenet)
print('Loading base network...')
net.base.load_state_dict(base_weights)
print('Base weights load done.')
# def xavier(param):
# init.xavier_uniform(param)
#
# def weights_init(m):
# for key in m.state_dict():
# if key.split('.')[-1] == 'weight':
# if 'conv' in key:
# init.kaiming_normal_(m.state_dict()[key], mode='fan_out')
# if 'bn' in key:
# m.state_dict()[key][...] = 1
# elif key.split('.')[-1] == 'bias':
# m.state_dict()[key][...] = 0
net.to(device)
optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
criterion = MultiBoxLoss(num_classes, 0.5, True, 0, True, 3, 0.5, False).to(device)
priorbox = PriorBox(cfg)
with torch.no_grad():
priors = priorbox.forward()
priors.to(device)
net.train()
# loss counters
loc_loss = 0 # epoch
conf_loss = 0
epoch = 0 + args.resume_epoch
epoch_size = len(dataset) // args.batch_size
max_iter = args.max_epoch * epoch_size
stepvalues_VOC = (150 * epoch_size, 200 * epoch_size, 250 * epoch_size)
stepvalues_COCO = (90 * epoch_size, 120 * epoch_size, 140 * epoch_size)
stepvalues = (stepvalues_VOC, stepvalues_COCO)[args.dataset == 'COCO']
print('Training', args.version, 'on', dataset.name)
step_index = 0
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
lr = args.lr
try:
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(data.DataLoader(dataset, batch_size,
shuffle=True, num_workers=args.num_workers,
collate_fn=detection_collate))
loc_loss = 0
conf_loss = 0
if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 == 0 and epoch > 200):
torch.save(net.state_dict(), args.save_folder + args.version + '_' + args.dataset + '_epoches_' +
repr(epoch) + '.pth')
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(args, optimizer, args.gamma, epoch, step_index, iteration, epoch_size)
# load train data
images, targets = next(batch_iterator)
images = Variable(images.to(device))
targets = [Variable(anno.to(device)) for anno in targets]
print(images)
print(targets)
# forward
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, priors, targets)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
t1 = time.time()
loc_loss += loss_l.item()
conf_loss += loss_c.item()
load_t1 = time.time()
if iteration % 10 == 0:
print('Epoch:' + repr(epoch) + ' || epochiter: ' + repr(iteration % epoch_size) + '/' + repr(epoch_size)
+ '|| Totel iter ' +
repr(iteration) + ' || L: %.4f C: %.4f||' % (
loss_l.item(), loss_c.item()) +
'Batch time: %.4f sec. ||' % (load_t1 - load_t0) + 'LR: %.8f' % (lr))
except KeyboardInterrupt:
print('Interrupted, try saving model...')
torch.save(net.state_dict(), args.save_folder +
'Final_' + args.version + '_' + args.dataset + '.pth')
torch.save(net.state_dict(), args.save_folder +
'Final_' + args.version + '_' + args.dataset + '.pth')
from torch.autograd import Variable
import numpy as np
import cv2
from matplotlib import pyplot as plt
from data import VOCDetection, VOC_ROOT, VOCAnnotationTransform
from ssd import build_ssd
if torch.cuda.is_available():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
net = build_ssd('test', 300, 21) # initialize SSD
net.load_weights('weights/ssd300_mAP_77.43_v2.pth')
# here we specify year (07 or 12) and dataset ('test', 'val', 'train')
testset = VOCDetection(VOC_ROOT, [('2007', 'val')], None, VOCAnnotationTransform())
img_id = 345 # change here for image id
image, image_name = testset.pull_image(img_id)
print(image_name)
rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# View the sampled input image before transform
# plt.figure(figsize=(6,6))
# plt.imshow(rgb_image)
# plt.show()
def base_transform(image):
x = cv2.resize(image, (300, 300)).astype(np.float32)
x -= (104.0, 117.0, 123.0)
x = x.astype(np.float32)
x = x[:, :, ::-1].copy()
def train():
net.train()
epoch = 0 + args.resume_epoch
print('Loading Dataset...')
dataset = VOCDetection(training_dataset, preproc(img_dim, rgb_mean),
AnnotationTransform())
epoch_size = math.ceil(len(dataset) / batch_size)
max_iter = max_epoch * epoch_size
stepvalues = (200 * epoch_size, 250 * epoch_size)
step_index = 0
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(
data.DataLoader(dataset,
batch_size,
shuffle=True,
num_workers=num_workers,
collate_fn=detection_collate))
if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 == 0
and epoch > 200):
torch.save(
net.state_dict(),
save_folder + 'FaceBoxes_epoch_' + str(epoch) + '.pth')
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(optimizer, gamma, epoch, step_index,
iteration, epoch_size)
# load train data
images, targets = next(batch_iterator)
images = images.to(device)
targets = [anno.to(device) for anno in targets]
# forward
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, priors, targets)
loss = cfg['loc_weight'] * loss_l + loss_c
loss.backward()
optimizer.step()
load_t1 = time.time()
batch_time = load_t1 - load_t0
eta = int(batch_time * (max_iter - iteration))
print(
'Epoch:{}/{} || Epochiter: {}/{} || Iter: {}/{} || L: {:.4f} C: {:.4f} || LR: {:.8f} || Batchtime: {:.4f} s || ETA: {}'
.format(epoch, max_epoch, (iteration % epoch_size) + 1,
epoch_size, iteration + 1, max_iter, loss_l.item(),
loss_c.item(), lr, batch_time,
str(datetime.timedelta(seconds=eta))))
torch.save(net.state_dict(), save_folder + 'Final_FaceBoxes.pth')
def evaluate_detections(box_list, output_dir, dataset):
write_voc_results_file(box_list, dataset)
do_python_eval(output_dir)
if __name__ == '__main__':
# load net
num_classes = len(labelmap) + 1 # +1 for background
net = build_refinedet('test', int(args.input_size),
num_classes) # initialize SSD
net.load_state_dict(torch.load(args.trained_model))
net.eval()
print('Finished loading model!')
# load data
dataset = VOCDetection(args.voc_root, [('2007', set_type)],
BaseTransform(int(args.input_size), dataset_mean),
VOCAnnotationTransform())
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
# evaluation
test_net(args.save_folder,
net,
args.cuda,
dataset,
BaseTransform(net.size, dataset_mean),
args.top_k,
int(args.input_size),
thresh=args.confidence_threshold)
else:
print('Unkown version!')
priorbox = PriorBox(cfg)
with torch.no_grad():
priors = priorbox.forward()
if args.cuda:
priors = priors.cuda()
# define test data set
img_size = (300, 512)[args.size == '512']
rgb_mean = (104, 117, 123)
img_set = ('SDYY', 'test')
transform = ImageAugment(img_size, rgb_mean)
testset = VOCDetection(VOCroot, img_set, CLASSES, transform)
num_classes = len(testset.classes)
# define net and load weight
net = build_net('test', img_size, num_classes)
state_dict = torch.load(args.trained_model)
new_state_dict = OrderedDict()
for k, v in state_dict.items():
head = k[:7]
if head == 'module.':
name = k[7:] # remove `module.`
else:
name = k
new_state_dict[name] = v
# load new state dict
def train():
net.train()
# loss counters
loc_loss = 0 # epoch
conf_loss = 0
epoch = 0 + args.resume_epoch
print('Loading Dataset...')
if args.dataset == 'VOC':
dataset = VOCDetection(VOCroot, train_sets,
preproc(img_dim, rgb_means, p),
AnnotationTransform())
elif args.dataset == 'COCO':
dataset = COCODetection(COCOroot, train_sets,
preproc(img_dim, rgb_means, p))
else:
print('Only VOC and COCO are supported now!')
return
epoch_size = len(dataset) // args.batch_size
max_iter = args.max_epoch * epoch_size
stepvalues_VOC = (150 * epoch_size, 200 * epoch_size, 250 * epoch_size)
stepvalues_COCO = (90 * epoch_size, 120 * epoch_size, 140 * epoch_size)
stepvalues = (stepvalues_VOC, stepvalues_COCO)[args.dataset == 'COCO']
print('Training', args.version, 'on', dataset.name)
step_index = 0
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
lr = args.lr
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(
data.DataLoader(dataset,
batch_size,
shuffle=True,
num_workers=args.num_workers,
collate_fn=detection_collate))
loc_loss = 0
conf_loss = 0
if (epoch % 10 == 0 and epoch > 0) or (epoch % 5 == 0
and epoch > 200):
torch.save(
net.state_dict(),
'delete10_11_' + args.save_folder + args.version + '_' +
args.dataset + '_epoches_' + repr(epoch) + '.pth')
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
lr = adjust_learning_rate(optimizer, args.gamma, epoch, step_index,
iteration, epoch_size)
# load train data
images, targets = next(batch_iterator)
#print(np.sum([torch.sum(anno[:,-1] == 2) for anno in targets]))
if args.cuda:
images = Variable(images.cuda())
targets = [Variable(anno.cuda()) for anno in targets]
else:
images = Variable(images)
targets = [Variable(anno) for anno in targets]
# forward
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, priors, targets)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
t1 = time.time()
loc_loss += loss_l.item()
conf_loss += loss_c.item()
load_t1 = time.time()
if iteration % 10 == 0:
print('Epoch:' + repr(epoch) + ' || epochiter: ' +
repr(iteration % epoch_size) + '/' + repr(epoch_size) +
'|| Totel iter ' + repr(iteration) +
' || L: %.4f C: %.4f||' % (loss_l.item(), loss_c.item()) +
'Batch time: %.4f sec. ||' % (load_t1 - load_t0) +
'LR: %.8f' % (lr))
torch.save(
net.state_dict(), args.save_folder + 'Final_' + args.version + '_' +
args.dataset + '.pth')
def test():
# get device
if args.cuda:
print('use cuda')
cudnn.benchmark = True
device = torch.device("cuda")
else:
device = torch.device("cpu")
# load net
num_classes = 80
if args.dataset == 'COCO':
cfg = config.coco_ab
testset = COCODataset(data_dir=args.dataset_root,
json_file='instances_val2017.json',
name='val2017',
img_size=cfg['min_dim'][0],
debug=args.debug)
elif args.dataset == 'VOC':
cfg = config.voc_ab
testset = VOCDetection(VOC_ROOT, [('2007', 'test')], None,
VOCAnnotationTransform())
if args.version == 'yolo_v2':
from models.yolo_v2 import myYOLOv2
net = myYOLOv2(device,
input_size=cfg['min_dim'],
num_classes=num_classes,
anchor_size=config.ANCHOR_SIZE_COCO)
print('Let us test yolo-v2 on the MSCOCO dataset ......')
elif args.version == 'yolo_v3':
from models.yolo_v3 import myYOLOv3
net = myYOLOv3(device,
input_size=cfg['min_dim'],
num_classes=num_classes,
anchor_size=config.MULTI_ANCHOR_SIZE_COCO)
elif args.version == 'tiny_yolo_v2':
from models.tiny_yolo_v2 import YOLOv2tiny
net = YOLOv2tiny(device,
input_size=cfg['min_dim'],
num_classes=num_classes,
anchor_size=config.ANCHOR_SIZE_COCO)
elif args.version == 'tiny_yolo_v3':
from models.tiny_yolo_v3 import YOLOv3tiny
net = YOLOv3tiny(device,
input_size=cfg['min_dim'],
num_classes=num_classes,
anchor_size=config.MULTI_ANCHOR_SIZE_COCO)
net.load_state_dict(torch.load(args.trained_model, map_location='cuda'))
net.to(device).eval()
print('Finished loading model!')
# evaluation
test_net(net,
device,
testset,
BaseTransform(net.input_size,
mean=(0.406, 0.456, 0.485),
std=(0.225, 0.224, 0.229)),
thresh=args.visual_threshold)
backbone = 'd-53'
yolo_net = YOLOv3SPP(device,
input_size=input_size,
num_classes=num_classes,
anchor_size=anchor_size,
backbone=backbone,
diou_nms=args.diou_nms)
print('Let us test yolo-v3-spp on the VOC dataset ......')
else:
print('Unknown version !!!')
exit()
# load net
yolo_net.load_state_dict(
torch.load(args.trained_model, map_location='cuda'))
yolo_net.eval()
print('Finished loading model!')
# load data
dataset = VOCDetection(
args.voc_root, [('2007', set_type)],
BaseTransform(yolo_net.input_size,
mean=(0.406, 0.456, 0.485),
std=(0.225, 0.224, 0.229)), VOCAnnotationTransform())
yolo_net = yolo_net.to(device)
# evaluation
with torch.no_grad():
test_net(yolo_net, dataset, device, input_size)
def train():
net.train()
# loss counters
loc_loss = 0 # epoch
conf_loss = 0
epoch = 0 + args.resume_epoch
print('Loading Dataset...')
if args.dataset == 'VOC':
dataset = VOCDetection(VOCroot, train_sets,
preproc(img_dim, rgb_means, p),
AnnotationTransform())
elif args.dataset == 'COCO':
dataset = COCODetection(COCOroot, train_sets,
preproc(img_dim, rgb_means, p))
else:
print('Only VOC and COCO are supported now!')
return
epoch_size = len(dataset) // args.batch_size
max_iter = args.max_epoch * epoch_size
stepvalues_VOC = (150 * epoch_size, 200 * epoch_size, 250 * epoch_size)
stepvalues_COCO = (90 * epoch_size, 120 * epoch_size, 140 * epoch_size)
stepvalues = (stepvalues_VOC, stepvalues_COCO)[args.dataset == 'COCO']
print('Training', args.version, 'on', dataset.name)
step_index = 0
if args.visdom:
# initialize visdom loss plot
lot = viz.line(X=torch.zeros((1, )).cpu(),
Y=torch.zeros((1, 3)).cpu(),
opts=dict(xlabel='Iteration',
ylabel='Loss',
title='Current SSD Training Loss',
legend=['Loc Loss', 'Conf Loss', 'Loss']))
epoch_lot = viz.line(X=torch.zeros((1, )).cpu(),
Y=torch.zeros((1, 3)).cpu(),
opts=dict(
xlabel='Epoch',
ylabel='Loss',
title='Epoch SSD Training Loss',
legend=['Loc Loss', 'Conf Loss', 'Loss']))
if args.resume_epoch > 0:
start_iter = args.resume_epoch * epoch_size
else:
start_iter = 0
lr = args.lr
log_file = open(log_file_path, 'a')
for iteration in range(start_iter, max_iter):
if iteration % epoch_size == 0:
# create batch iterator
batch_iterator = iter(
data.DataLoader(dataset,
batch_size,
shuffle=True,
num_workers=args.num_workers,
collate_fn=detection_collate))
loc_loss = 0
conf_loss = 0
if epoch % args.save_frequency == 0 and epoch > 0:
torch.save(
net.state_dict(), save_folder + args.version + '_' +
args.dataset + '_epoches_' + repr(epoch) + '.pth')
epoch += 1
load_t0 = time.time()
if iteration in stepvalues:
step_index += 1
if args.visdom:
viz.line(
X=torch.ones((1, 3)).cpu() * epoch,
Y=torch.Tensor([loc_loss, conf_loss, loc_loss + conf_loss
]).unsqueeze(0).cpu() / epoch_size,
win=epoch_lot,
update='append')
lr = adjust_learning_rate(optimizer, args.gamma, epoch, step_index,
iteration, epoch_size)
# load train data
images, targets = next(batch_iterator)
#print(np.sum([torch.sum(anno[:,-1] == 2) for anno in targets]))
if args.cuda:
images = Variable(images.cuda())
targets = [
Variable(anno.cuda(), volatile=True) for anno in targets
]
else:
images = Variable(images)
targets = [Variable(anno, volatile=True) for anno in targets]
# forward
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, priors, targets)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
t1 = time.time()
loc_loss += loss_l.data[0]
conf_loss += loss_c.data[0]
load_t1 = time.time()
if iteration % 10 == 0:
print('Epoch:' + repr(epoch) + ' || epochiter: ' +
repr(iteration % epoch_size) + '/' + repr(epoch_size) +
'|| Totel iter ' + repr(iteration) +
' || L: %.4f C: %.4f||' % (loss_l.data[0], loss_c.data[0]) +
'Batch time: %.4f sec. ||' % (load_t1 - load_t0) +
'LR: %.8f' % (lr))
log_file.write('Epoch:' + repr(epoch) + ' || epochiter: ' +
repr(iteration % epoch_size) + '/' +
repr(epoch_size) + '|| Totel iter ' +
repr(iteration) + ' || L: %.4f C: %.4f||' %
(loss_l.data[0], loss_c.data[0]) +
'Batch time: %.4f sec. ||' % (load_t1 - load_t0) +
'LR: %.8f' % (lr) + '\n')
if args.visdom and args.send_images_to_visdom:
random_batch_index = np.random.randint(images.size(0))
viz.image(images.data[random_batch_index].cpu().numpy())
if args.visdom:
viz.line(X=torch.ones((1, 3)).cpu() * iteration,
Y=torch.Tensor([
loss_l.data[0], loss_c.data[0],
loss_l.data[0] + loss_c.data[0]
]).unsqueeze(0).cpu(),
win=lot,
update='append')
if iteration == 0:
viz.line(X=torch.zeros((1, 3)).cpu(),
Y=torch.Tensor(
[loc_loss, conf_loss,
loc_loss + conf_loss]).unsqueeze(0).cpu(),
win=epoch_lot,
update=True)
log_file.close()
torch.save(
net.state_dict(),
save_folder + 'Final_' + args.version + '_' + args.dataset + '.pth')
def train():
net.train()
# loss counters
loc_loss = 0 # epoch
conf_loss = 0
epoch = 0
print('Loading Dataset...')
dataset = VOCDetection(args.voc_root, train_sets, SSDAugmentation(
ssd_dim, means), AnnotationTransform())
epoch_size = len(dataset) // args.batch_size
print('Training SSD on', dataset.name)
step_index = 0
if args.visdom:
# initialize visdom loss plot
lot = viz.line(
X=torch.zeros((1,)).cpu(),
Y=torch.zeros((1, 3)).cpu(),
opts=dict(
xlabel='Iteration',
ylabel='Loss',
title='Current SSD Training Loss',
legend=['Loc Loss', 'Conf Loss', 'Loss']
)
)
epoch_lot = viz.line(
X=torch.zeros((1,)).cpu(),
Y=torch.zeros((1, 3)).cpu(),
opts=dict(
xlabel='Epoch',
ylabel='Loss',
title='Epoch SSD Training Loss',
legend=['Loc Loss', 'Conf Loss', 'Loss']
)
)
batch_iterator = None
data_loader = data.DataLoader(dataset, batch_size, num_workers=args.num_workers,
shuffle=True, collate_fn=detection_collate, pin_memory=True)
for iteration in range(args.start_iter, max_iter):
if (not batch_iterator) or (iteration % epoch_size == 0):
# create batch iterator
batch_iterator = iter(data_loader)
if iteration in stepvalues:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
if args.visdom:
viz.line(
X=torch.ones((1, 3)).cpu() * epoch,
Y=torch.Tensor([loc_loss, conf_loss,
loc_loss + conf_loss]).unsqueeze(0).cpu() / epoch_size,
win=epoch_lot,
update='append'
)
# reset epoch loss counters
loc_loss = 0
conf_loss = 0
epoch += 1
# load train data
images, targets = next(batch_iterator)
#print('--', images[0].size())
if args.cuda:
images = Variable(images.cuda())
targets = [Variable(anno.cuda(), volatile=True) for anno in targets]
else:
images = Variable(images)
targets = [Variable(anno, volatile=True) for anno in targets]
# forward
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, targets)
loss = loss_l + loss_c
loss.backward()
optimizer.step()
t1 = time.time()
loc_loss += loss_l.data[0]
conf_loss += loss_c.data[0]
if iteration % 10 == 0:
print('Timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.data[0]), end=' ')
if args.visdom and args.send_images_to_visdom:
random_batch_index = np.random.randint(images.size(0))
viz.image(images.data[random_batch_index].cpu().numpy())
if args.visdom:
viz.line(
X=torch.ones((1, 3)).cpu() * iteration,
Y=torch.Tensor([loss_l.data[0], loss_c.data[0],
loss_l.data[0] + loss_c.data[0]]).unsqueeze(0).cpu(),
win=lot,
update='append'
)
# hacky fencepost solution for 0th epoch plot
if iteration == 0:
viz.line(
X=torch.zeros((1, 3)).cpu(),
Y=torch.Tensor([loc_loss, conf_loss,
loc_loss + conf_loss]).unsqueeze(0).cpu(),
win=epoch_lot,
update=True
)
if iteration % 5000 == 0:
print('Saving state, iter:', iteration)
torch.save(ssd_net.state_dict(), '/media/maxiaoyu/data/checkpoint/weights/ssd300_0712_' + repr(iteration) + '.pth')
#print('finish one iterate')
# used for eval
torch.save(ssd_net.state_dict(), args.save_folder + '' + args.version + '.pth')
def do_test(args, model, detector, max_per_image=200, thresh=0.01):
if args.dataset == 'VOC':
dataset = VOCDetection(
args, VOCroot, [('2007', 'test')], None,
AnnotationTransform(0 if args.setting ==
'transfer' else args.split), True)
elif args.dataset == 'COCO':
dataset = COCODetection(COCOroot, [('2014', 'split_nonvoc_minival')],
None)
else:
raise ValueError(f"Unknown dataset: {args.dataset}")
num_images = len(dataset)
all_boxes = [[[] for _ in range(num_images)] for _ in range(num_classes)]
transform = BaseTransform(model.size, rgb_means, (2, 0, 1))
_t = {'im_detect': Timer(), 'misc': Timer()}
det_file = os.path.join(args.save_folder, 'detections.pkl')
if args.retest:
f = open(det_file, 'rb')
all_boxes = pickle.load(f)
logger.info('Evaluating detections')
dataset.evaluate_detections(all_boxes, args.save_folder)
return
for i in range(num_images):
img = dataset.pull_image(i)
scale = torch.Tensor(
[img.shape[1], img.shape[0], img.shape[1],
img.shape[0]]).to(model.device)
with torch.no_grad():
x = transform(img).unsqueeze(0)
_t['im_detect'].tic()
pred = model(x) # forward pass
boxes, scores = detector.forward(pred, priors)
detect_time = _t['im_detect'].toc()
boxes = boxes[0] # percent and point form detection boxes
scores = scores[0] # [1, num_priors, num_classes]
boxes *= scale # scale each detection back up to the image
boxes = boxes.cpu().numpy()
scores = scores.cpu().numpy()
_t['misc'].tic()
for j in range(1, num_classes):
inds = np.where(scores[:, j] > thresh)[0]
if len(inds) == 0:
all_boxes[j][i] = np.empty([0, 5], dtype=np.float32)
continue
c_bboxes = boxes[inds]
c_scores = scores[inds, j]
c_dets = np.hstack(
(c_bboxes, c_scores[:, np.newaxis])).astype(np.float32,
copy=False)
keep = nms(c_dets, 0.45, force_cpu=args.cpu)
c_dets = c_dets[keep, :]
all_boxes[j][i] = c_dets
if max_per_image > 0:
image_scores = np.hstack(
[all_boxes[j][i][:, -1] for j in range(1, num_classes)])
if len(image_scores) > max_per_image:
image_thresh = np.sort(image_scores)[-max_per_image]
for j in range(1, num_classes):
keep = np.where(all_boxes[j][i][:, -1] >= image_thresh)[0]
all_boxes[j][i] = all_boxes[j][i][keep, :]
nms_time = _t['misc'].toc()
if i % 20 == 0:
logger.info('im_detect: {:d}/{:d} {:.3f}s {:.3f}s'.format(
i + 1, num_images, detect_time, nms_time))
_t['im_detect'].clear()
_t['misc'].clear()
with open(det_file, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
logger.info('Evaluating detections')
dataset.evaluate_detections(all_boxes, args.save_folder)
