def train():
if args.dataset == 'COCO':
cfg = coco_refinedet[args.input_size]
root = '/home/soyeol/dataset/coco2014'
dataset = COCODetection(root=root,
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
print("len_dataset(COCO) :", len(dataset))
elif args.dataset == 'VOC':
cfg = voc_refinedet[args.input_size]
root = '/home/soyeol/dataset/VOCdevkit'
# root = "/content/drive/MyDrive/Colab Notebooks/dataset/VOCdevkit"
# val_dataset = VOCDetection(root=root, transform=BaseTransform(args.input_size, mean=(104, 117, 123)),
# image_sets=[('2007', 'test')], mode='test')
val_dataset = infer.InferDataset(
os.path.join(root, 'VOC2007', 'JPEGImages'))
print("len_dataset(VOC) :", len(val_dataset))
net = build_refinedet('train', cfg['min_dim'],
cfg['num_classes']) # , batch_size=args.batch_size)
print(net)
torch.multiprocessing.set_start_method('spawn')
if args.cuda:
# net = torch.nn.DataParallel(refinedet_net)
cudnn.benchmark = True
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
state = torch.load(args.resume)
net.load_state_dict(state['model'])
if args.cuda:
net = net.cuda()
iteration = state['iter']
print(args)
net.phase = 'test'
map = _infer(net, args, iteration, val_dataset, cfg) # dataset[1])
net.phase = 'train'
print('mAP:', map)
# from models.refinedet_test import build_refinedet
import numpy as np
import cv2
import os
def check_dir(dir):
if not os.path.exists(dir):
os.makedirs(dir)
if __name__ == '__main__':
torch.set_default_tensor_type('torch.cuda.FloatTensor')
num_classes = 2# len(labelmap) + 1 # +1 for background
net = build_refinedet(320, num_classes) # initialize SSD
net.load_state_dict(torch.load("/home/pengyang/Code/gitlab/follow_model_tx2/pytorch_model_pth/resnet50-86_refinedet/RefineDet320_VOC_res86_body_head.pth"))
net.eval()
net = net.cuda()
cudnn.benchmark = True
print('Finished loading model!')
VOC_ROOT = "/home/py/Disk700G/2019code/Train_data/VOCdevkit/"
src = os.path.join(VOC_ROOT, "VOC2007/ImageSets/Main")
imgPath = os.path.join(VOC_ROOT, "VOC2007/JPEGImages")
dst = "/home/lzm/Disk2/work_dl/Pytorch_refinedet/Data_dir/Test_model/RefineDet320_VOC_140000_withhead"
check_dir(dst)
for dirs in os.listdir(src):
def train():
if args.dataset == 'COCO':
'''if args.dataset_root == VOC_ROOT:
if not os.path.exists(COCO_ROOT):
parser.error('Must specify dataset_root if specifying dataset')
print("WARNING: Using default COCO dataset_root because " +
"--dataset_root was not specified.")
args.dataset_root = COCO_ROOT
cfg = coco
dataset = COCODetection(root=args.dataset_root,
transform=SSDAugmentation(cfg['min_dim'],
MEANS))'''
elif args.dataset == 'VOC':
'''if args.dataset_root == COCO_ROOT:
parser.error('Must specify dataset if specifying dataset_root')'''
cfg = voc_refinedet[args.input_size]
dataset = VOCDetection(root=args.dataset_root,
transform=SSDAugmentation(cfg['min_dim'],
MEANS))
if args.visdom:
import visdom
viz = visdom.Visdom()
refinedet_net = build_refinedet('train', cfg['min_dim'], cfg['num_classes'])
net = refinedet_net
print(net)
#input()
if args.cuda:
net = torch.nn.DataParallel(refinedet_net)
cudnn.benchmark = True
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
refinedet_net.load_weights(args.resume)
else:
#vgg_weights = torch.load(args.save_folder + args.basenet)
vgg_weights = torch.load(args.basenet)
print('Loading base network...')
refinedet_net.vgg.load_state_dict(vgg_weights)
if args.cuda:
net = net.cuda()
if not args.resume:
print('Initializing weights...')
# initialize newly added layers' weights with xavier method
refinedet_net.extras.apply(weights_init)
refinedet_net.arm_loc.apply(weights_init)
refinedet_net.arm_conf.apply(weights_init)
refinedet_net.odm_loc.apply(weights_init)
refinedet_net.odm_conf.apply(weights_init)
#refinedet_net.tcb.apply(weights_init)
refinedet_net.tcb0.apply(weights_init)
refinedet_net.tcb1.apply(weights_init)
refinedet_net.tcb2.apply(weights_init)
optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum,
weight_decay=args.weight_decay)
arm_criterion = RefineDetMultiBoxLoss(2, 0.5, True, 0, True, 3, 0.5,
False, args.cuda)
odm_criterion = RefineDetMultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5,
False, args.cuda, use_ARM=True)
net.train()
# loss counters
arm_loc_loss = 0
arm_conf_loss = 0
odm_loc_loss = 0
odm_conf_loss = 0
epoch = 0
print('Loading the dataset...')
epoch_size = len(dataset) // args.batch_size
print('Training RefineDet on:', dataset.name)
print('Using the specified args:')
print(args)
step_index = 0
if args.visdom:
vis_title = 'RefineDet.PyTorch on ' + dataset.name
vis_legend = ['Loc Loss', 'Conf Loss', 'Total Loss']
iter_plot = create_vis_plot('Iteration', 'Loss', vis_title, vis_legend)
epoch_plot = create_vis_plot('Epoch', 'Loss', vis_title, vis_legend)
data_loader = data.DataLoader(dataset, args.batch_size,
num_workers=args.num_workers,
shuffle=True, collate_fn=detection_collate,
pin_memory=True)
print("data loader len: {}".format(len(data_loader)))
# create batch iterator
batch_iterator = iter(data_loader)
for iteration in range(args.start_iter, cfg['max_iter']):
if args.visdom and iteration != 0 and (iteration % epoch_size == 0):
update_vis_plot(epoch, arm_loc_loss, arm_conf_loss, epoch_plot, None,
'append', epoch_size)
# reset epoch loss counters
arm_loc_loss = 0
arm_conf_loss = 0
odm_loc_loss = 0
odm_conf_loss = 0
epoch += 1
if iteration in cfg['lr_steps']:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
# load train data
try:
images, targets = next(batch_iterator)
except StopIteration:
batch_iterator = iter(data_loader)
images, targets = next(batch_iterator)
if args.cuda:
images = images.cuda()
targets = [ann.cuda() for ann in targets]
else:
images = images
targets = [ann for ann in targets]
# forward
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
arm_loss_l, arm_loss_c = arm_criterion(out, targets)
odm_loss_l, odm_loss_c = odm_criterion(out, targets)
#input()
arm_loss = arm_loss_l + arm_loss_c
odm_loss = odm_loss_l + odm_loss_c
loss = arm_loss + odm_loss
loss.backward()
optimizer.step()
t1 = time.time()
arm_loc_loss += arm_loss_l.item()
arm_conf_loss += arm_loss_c.item()
odm_loc_loss += odm_loss_l.item()
odm_conf_loss += odm_loss_c.item()
if iteration % 100 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || ARM_L Loss: %.4f ARM_C Loss: %.4f ODM_L Loss: %.4f ODM_C Loss: %.4f total loss: %.4f||' \
% (arm_loss_l.item(), arm_loss_c.item(), odm_loss_l.item(), odm_loss_c.item(), loss.item()), end=' ')
if args.visdom:
update_vis_plot(iteration, arm_loss_l.data[0], arm_loss_c.data[0],
iter_plot, epoch_plot, 'append')
if iteration != 0 and iteration % 5000 == 0:
print('Saving state, iter:', iteration)
torch.save(refinedet_net.state_dict(), args.save_folder
+ '/RefineDet{}_{}_{}.pth'.format(args.input_size, args.dataset,
repr(iteration)))
torch.save(refinedet_net.state_dict(), args.save_folder
+ '/RefineDet{}_{}_final.pth'.format(args.input_size, args.dataset))
def train():
if args.dataset == 'VOC':
'''if args.dataset_root == COCO_ROOT:
parser.error('Must specify dataset if specifying dataset_root')'''
cfg = voc_refinedet[args.input_size]
dataset = VOCDetection(root=args.dataset_root,
transform=SSDAugmentation(cfg['min_dim'],
MEANS))
#可视化
viz = Visdom(server='http://127.0.0.1', port=8097)
assert viz.check_connection()
viz.line([[0.0, 0.0]], [0.], win='ARM_LOSS', opts=dict(title='arm_loss', legend=['ARM_L Loss', 'ARM_C Loss']))
viz.line([[0.0, 0.0]], [0.], win='ODM_LOSS', opts=dict(title='omd_loss', legend=['ODM_L Loss', 'ODM_C Loss']))
refinedet_net = build_refinedet(cfg['min_dim'], cfg['num_classes'])
net = refinedet_net
print(net)
#input()
# if args.cuda:
# net = torch.nn.DataParallel(refinedet_net)
# cudnn.benchmark = True
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
refinedet_net.load_weights(args.resume)
# else:
# refinedet_net.load_state_dict(torch.load(args.basenet), False)
if args.cuda:
net = net.cuda()
if not args.resume:
print('Initializing weights...')
# initialize newly added layers' weights with xavier method
refinedet_net.extras.apply(weights_init)
refinedet_net.arm_loc.apply(weights_init)
refinedet_net.arm_conf.apply(weights_init)
refinedet_net.odm_loc.apply(weights_init)
refinedet_net.odm_conf.apply(weights_init)
#refinedet_net.tcb.apply(weights_init)
refinedet_net.tcb0.apply(weights_init)
refinedet_net.tcb1.apply(weights_init)
refinedet_net.tcb2.apply(weights_init)
optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum,weight_decay=args.weight_decay)
# optimizer = torch.optim.Adam(net.parameters(), lr=args.lr)
arm_criterion = RefineDetMultiBoxLoss(2, 0.5, True, 0, True, 3, 0.5,
False, args.cuda)
odm_criterion = RefineDetMultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5,
False, args.cuda, use_ARM=True)
net.train()
# loss counters
arm_loc_loss = 0
arm_conf_loss = 0
odm_loc_loss = 0
odm_conf_loss = 0
epoch = 0
print('Loading the dataset...')
epoch_size = len(dataset) // args.batch_size
print('Training RefineDet on:', dataset.name)
print('Using the specified args:')
print(args)
step_index = 0
data_loader = data.DataLoader(dataset, args.batch_size,
num_workers=args.num_workers,
shuffle=True, collate_fn=detection_collate,
pin_memory=True)
# create batch iterator
batch_iterator = iter(data_loader)
for iteration in range(args.start_iter, cfg['max_iter']):
if iteration in cfg['lr_steps']:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
# load train data
try:
images, targets = next(batch_iterator)
except StopIteration:
batch_iterator = iter(data_loader)
images, targets = next(batch_iterator)
if args.cuda:
images = images.cuda()
targets = [ann.cuda() for ann in targets]
else:
images = images
targets = [ann for ann in targets]
# forward
t0 = time.time()
out = net(images)
# print("out:", len(out))
# backprop
optimizer.zero_grad()
arm_loss_l, arm_loss_c = arm_criterion(out, targets)
odm_loss_l, odm_loss_c = odm_criterion(out, targets)
#input()
arm_loss = arm_loss_l + arm_loss_c
odm_loss = odm_loss_l + odm_loss_c
loss = arm_loss + odm_loss
loss.backward()
optimizer.step()
t1 = time.time()
arm_loc_loss += arm_loss_l.item()
arm_conf_loss += arm_loss_c.item()
odm_loc_loss += odm_loss_l.item()
odm_conf_loss += odm_loss_c.item()
if iteration % 10 == 0:
# viz.line()
viz.line([[arm_loss_l.item(), arm_loss_c.item()]], [iteration], win='ARM_LOSS', update='append')
viz.line([[odm_loss_l.item(), odm_loss_c.item()]], [iteration], win='ODM_LOSS', update='append')
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' || ARM_L Loss: %.4f ARM_C Loss: %.4f ODM_L Loss: %.4f ODM_C Loss: %.4f ||' \
% (arm_loss_l.item(), arm_loss_c.item(), odm_loss_l.item(), odm_loss_c.item()), end=' ')
if iteration != 0 and iteration % 2000== 0:
print('Saving state, iter:', iteration)
torch.save(refinedet_net.state_dict(), args.save_folder
+ '/RefineDet{}_{}_{}.pth'.format(args.input_size, args.dataset,
repr(iteration)))
torch.save(refinedet_net.state_dict(), args.save_folder
+ '/RefineDet{}_{}_final.pth'.format(args.input_size, args.dataset))
def train():
if args.dataset == 'COCO':
'''if args.dataset_root == VOC_ROOT:
if not os.path.exists(COCO_ROOT):
parser.error('Must specify dataset_root if specifying dataset')
print("WARNING: Using default COCO dataset_root because " +
"--dataset_root was not specified.")
args.dataset_root = COCO_ROOT
cfg = coco
dataset = COCODetection(root=args.dataset_root,
transform=SSDAugmentation(cfg['min_dim'],
MEANS))'''
elif args.dataset == 'VOC':
'''if args.dataset_root == COCO_ROOT:
parser.error('Must specify dataset if specifying dataset_root')'''
cfg = voc_refinedet[args.input_size]
dataset = VOCDetection(root=args.dataset_root,
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
if args.visdom:
import visdom
viz = visdom.Visdom()
if args.addCbam:
print("add CBAM")
refinedet_net = build_refinedet_cbam('train', cfg['min_dim'],
cfg['num_classes'])
elif args.addSam:
print("add SAM")
refinedet_net = build_refinedet_sam('train', cfg['min_dim'],
cfg['num_classes'])
elif args.addSE:
print("add SE")
refinedet_net = build_refinedet_se('train', cfg['min_dim'],
cfg['num_classes'])
elif args.addBA_TCB:
print("add BA-TCB")
refinedet_net = build_refinedet_novel('train', cfg['min_dim'],
cfg['num_classes'])
elif args.addBA_TCB_CAM_Mish:
print("add BA-TCB_cam_mish")
refinedet_net = build_refinedet_novel_cam_mish('train', cfg['min_dim'],
cfg['num_classes'])
elif args.addBA_TCB_CAM_Swish:
print("add BA-TCB_cam_swish")
refinedet_net = build_refinedet_novel_cam_swish(
'train', cfg['min_dim'], cfg['num_classes'])
elif args.addBA_TCB_CAM:
print("add BA-TCB_cam")
refinedet_net = build_refinedet_novel_cam('train', cfg['min_dim'],
cfg['num_classes'])
elif args.addBA_TCB_SAM:
print("add BA-TCB_sam")
refinedet_net = build_refinedet_novel_sam('train', cfg['min_dim'],
cfg['num_classes'])
elif args.addCam:
print("add CAM")
refinedet_net = build_refinedet_cam('train', cfg['min_dim'],
cfg['num_classes'])
elif args.useMish:
print("using Mish activity")
refinedet_net = build_refinedet_mish('train', cfg['min_dim'],
cfg['num_classes'])
elif args.useSwish:
print("using Swish activity")
refinedet_net = build_refinedet_swish('train', cfg['min_dim'],
cfg['num_classes'])
else:
print("using ReLU")
refinedet_net = build_refinedet('train', cfg['min_dim'],
cfg['num_classes'])
net = refinedet_net
print(net)
#input()
if args.cuda:
net = torch.nn.DataParallel(refinedet_net)
cudnn.benchmark = True
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
refinedet_net.load_weights(args.resume)
else:
vgg_weights = torch.load(args.basenet)
print('Loading base network...')
refinedet_net.vgg.load_state_dict(vgg_weights)
if args.cuda:
net = net.cuda()
if not args.resume:
print('Initializing weights...')
# initialize newly added layers' weights with xavier method
refinedet_net.extras.apply(weights_init)
refinedet_net.arm_loc.apply(weights_init)
refinedet_net.arm_conf.apply(weights_init)
refinedet_net.odm_loc.apply(weights_init)
refinedet_net.odm_conf.apply(weights_init)
#refinedet_net.tcb.apply(weights_init)
refinedet_net.tcb0.apply(weights_init)
refinedet_net.tcb1.apply(weights_init)
refinedet_net.tcb2.apply(weights_init)
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.labelSmoothing:
print("Using Label Smoothing")
arm_criterion = RefineDetMultiBoxLoss1(2, 0.5, True, 0, True, 3, 0.5,
False, args.cuda)
odm_criterion = RefineDetMultiBoxLoss1(cfg['num_classes'],
0.5,
True,
0,
True,
3,
0.5,
False,
args.cuda,
use_ARM=True)
elif args.GIoULoss:
print("Using GIoULoss")
arm_criterion = RefineDetMultiBoxLoss2(2, 0.5, True, 0, True, 3, 0.5,
False, args.cuda)
odm_criterion = RefineDetMultiBoxLoss2(cfg['num_classes'],
0.5,
True,
0,
True,
3,
0.5,
False,
args.cuda,
use_ARM=True)
elif args.CIoULoss:
print("Using CIoULoss")
arm_criterion = RefineDetMultiBoxLoss4(2, 0.5, True, 0, True, 3, 0.5,
False, args.cuda)
odm_criterion = RefineDetMultiBoxLoss4(cfg['num_classes'],
0.5,
True,
0,
True,
3,
0.5,
False,
args.cuda,
use_ARM=True)
elif args.DIoULoss:
print("Using DIoULoss")
arm_criterion = RefineDetMultiBoxLoss3(2, 0.5, True, 0, True, 3, 0.5,
False, args.cuda)
odm_criterion = RefineDetMultiBoxLoss3(cfg['num_classes'],
0.5,
True,
0,
True,
3,
0.5,
False,
args.cuda,
use_ARM=True)
else:
print("Using L1 smooth loss")
arm_criterion = RefineDetMultiBoxLoss(2, 0.5, True, 0, True, 3, 0.5,
False, args.cuda)
odm_criterion = RefineDetMultiBoxLoss(cfg['num_classes'],
0.5,
True,
0,
True,
3,
0.5,
False,
args.cuda,
use_ARM=True)
net.train()
# loss counters
arm_loc_loss = 0
arm_conf_loss = 0
odm_loc_loss = 0
odm_conf_loss = 0
loss_ = 0
epoch = 0
best_loss = 100
print('Loading the dataset...')
epoch_size = len(dataset) // args.batch_size
print(len(dataset))
print('Training RefineDet on:', dataset.name)
print('Using the specified args:')
print(args)
step_index = 0
if args.visdom:
vis_title = 'RefineDet.PyTorch on ' + dataset.name
vis_legend = ['Loc Loss', 'Conf Loss', 'Total Loss']
iter_plot = create_vis_plot('Iteration', 'Loss', vis_title, vis_legend)
epoch_plot = create_vis_plot('Epoch', 'Loss', vis_title, vis_legend)
data_loader = data.DataLoader(dataset,
args.batch_size,
num_workers=args.num_workers,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
# create batch iterator
batch_iterator = iter(data_loader)
# lr_scheduler = CosineWarmupLr(optimizer, epoch_size, ,base_lr=args.lr,warmup_epochs=10)
scheduler = None
if args.cosineAnnealLR:
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
120000,
eta_min=0,
last_epoch=-1)
pltx = []
plty = []
pltz = []
pltm = []
for iteration in range(args.start_iter, cfg['max_iter']):
if args.visdom and iteration != 0 and (iteration % epoch_size == 0):
update_vis_plot(epoch, arm_loc_loss, arm_conf_loss, epoch_plot,
None, 'append', epoch_size)
# reset epoch loss counters
arm_loc_loss = 0
arm_conf_loss = 0
odm_loc_loss = 0
odm_conf_loss = 0
loss_ = 0
epoch += 1
if args.cosineAnnealLR:
pass
else:
if iteration in cfg['lr_steps']:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
# load train data
try:
images, targets = next(batch_iterator)
except StopIteration:
batch_iterator = iter(data_loader)
images, targets = next(batch_iterator)
if args.cuda:
images = images.cuda()
targets = [ann.cuda() for ann in targets]
else:
images = images
targets = [ann for ann in targets]
# forward
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
arm_loss_l, arm_loss_c = arm_criterion(out, targets)
odm_loss_l, odm_loss_c = odm_criterion(out, targets)
#input()
arm_loss = arm_loss_l + arm_loss_c
odm_loss = odm_loss_l + odm_loss_c
loss = arm_loss + odm_loss
loss.backward()
optimizer.step()
if args.cosineAnnealLR:
scheduler.step()
t1 = time.time()
loss_ += loss.item()
arm_loc_loss += arm_loss_l.item()
arm_conf_loss += arm_loss_c.item()
odm_loc_loss += odm_loss_l.item()
odm_conf_loss += odm_loss_c.item()
for param_group in optimizer.param_groups:
out_lr = param_group['lr']
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('iter ' + repr(iteration) + ' ||LR:%.8f Total Loss:%.4f ARM_L Loss: %.4f ARM_C Loss: %.4f ODM_L Loss: %.4f ODM_C Loss: %.4f ||' \
% (out_lr,loss.item(),arm_loss_l.item(), arm_loss_c.item(), odm_loss_l.item(), odm_loss_c.item()), end=' ')
if args.visdom:
update_vis_plot(iteration, arm_loss_l.data[0], arm_loss_c.data[0],
iter_plot, epoch_plot, 'append')
if iteration != 0 and iteration % 5000 == 0:
print('Saving state, iter:', iteration)
torch.save(
refinedet_net.state_dict(),
args.save_folder + '/RefineDet{}_voc0712_{}_.pth'.format(
args.input_size, repr(iteration)))
if iteration % 500 == 0:
pltx.append(iteration)
plty.append(loss.item())
pltz.append(arm_loss_c.item() + arm_loss_l + arm_loss_l.item())
pltm.append(odm_loss_c.item() + odm_loss_l.item())
# if loss<best_loss:
# best_loss=loss
# torch.save(refinedet_net.state_dict(), args.save_folder
# + '/RefineDet{}_{}_TILDA_sam_best_.pth'.format(args.input_size, args.dataset))
torch.save(
refinedet_net.state_dict(), args.save_folder +
'/RefineDet{}_voc0712_final120000.pth'.format(args.input_size))
plt.plot(pltx, plty, label='Total loss')
plt.plot(pltx, pltz, label='CA-ARM loss')
plt.plot(pltx, pltm, label='ODM loss')
plt.xlabel("iteraton")
plt.ylabel('total loss')
plt.legend()
plt.savefig('VOC loss.jpg')
import torch
import torch.nn as nn
import struct
from models.refinedet import build_refinedet
num_classes = 25
path_model = "/data_2/project_2021/pytorch_refinedet/2021/20210308.pth"
path_save_wts = "./refinedet0312.wts"
input_size = 320
net = build_refinedet('test', input_size, num_classes) # initialize net
net.load_state_dict(torch.load(path_model))
net.eval()
f = open(path_save_wts, 'w')
f.write('{}\n'.format(len(net.state_dict().keys())))
for k, v in net.state_dict().items():
vr = v.reshape(-1).cpu().numpy()
f.write('{} {} '.format(k, len(vr)))
for vv in vr:
f.write(' ')
f.write(struct.pack('>f', float(vv)).hex())
f.write('\n')
print("success generate wts!")
with open(det_file, 'wb') as f:
pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)
print('Evaluating detections')
evaluate_detections(all_boxes, output_dir, dataset)
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,
#from ssd import build_ssd
from models.refinedet import build_refinedet
if torch.cuda.is_available():
if args.cuda:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
if not args.cuda:
print(
"WARNING: It looks like you have a CUDA device, but aren't using \
CUDA. Run with --cuda for optimal eval speed.")
torch.set_default_tensor_type('torch.FloatTensor')
else:
torch.set_default_tensor_type('torch.FloatTensor')
#net = build_ssd('test', 300, 21) # initialize SSD
net = build_refinedet('test', 320, 4) # initialize SSD
net.load_state_dict(torch.load(args.weights))
transform = BaseTransform(net.size,
(104 / 256.0, 117 / 256.0, 123 / 256.0))
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
cv2_demo(net.eval(), transform)
# stop the timer and display FPS information
#fps.stop()
#print("[INFO] elasped time: {:.2f}".format(fps.elapsed()))
#print("[INFO] approx. FPS: {:.2f}".format(fps.fps()))
# cleanup
model,
cuda,
dataset,
BaseTransform(im_size, dataset_mean),
top_k,
im_size,
thresh=thresh)
return map
if __name__ == '__main__':
# load net
num_classes = len(labelmap) + 1 # +1 for background
image_size = 320
net = build_refinedet('test', image_size, num_classes)
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(),phase='valid')
if args.cuda:
net = net.cuda()
cudnn.benchmark = True
# # evaluation
# eval_net(args.save_folder, net, args.cuda, dataset,
# BaseTransform(net.size, dataset_mean), args.top_k, 300,
# thresh=args.confidence_threshold)
def __len__(self):
return len(self.files)
parser = argparse.ArgumentParser(description='Single Shot MultiBox Detection')
parser.add_argument('--trained_model', default='weights/RefineDet320_591.pth',type=str, help='Trained state_dict file path to open')
parser.add_argument('--save_folder', default='output/', type=str,help='Dir to save results')
parser.add_argument('--dataset_root', default='data/samples', help='Dataset root directory path')
parser.add_argument("--batch_size", type=int, default=1, help="size of the batches")
parser.add_argument("--n_cpu", type=int, default=0, help="number of cpu threads to use during batch generation")
parser.add_argument("--img_size", type=int, default=320, help="size of each image dimension")
args = parser.parse_args()
os.makedirs("refinedet_output", exist_ok=True)
net = build_refinedet('test', 320, 7)
net.load_weights(args.trained_model)
net = net.cuda()
dataloader = DataLoader(
ImageFolder(args.dataset_root, img_size=args.img_size),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.n_cpu,
)
TIME=0
for batch_i, (img_paths, input_imgs) in enumerate(dataloader):
# Configure input
# input_imgs = Variable(input_imgs.type(Tensor))
# print(img_paths[0])
# resnet = models.alexnet()
# resnet = models.resnet152()
model = model
foo(model)
input = Variable(torch.rand(3, size, size).unsqueeze(0),
requires_grad=True).cuda()
out = model(input)
total_flops = (sum(list_conv) + sum(list_linear) + sum(list_bn) +
sum(list_relu) + sum(list_pooling))
print(' + Number of FLOPs: %.2fG' % (total_flops / 1e9))
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Set up model
net1 = build_ssd('train', 300, 7).to(device)
net2 = build_refinedet('train', 320, 7).to(device)
# summary(net1, input_size=(3,300,300))
total1 = sum([param.nelement() for param in net1.parameters()])
total2 = sum([param.nelement() for param in net2.parameters()])
print(' + SSD Number of params: %.2fM' % (total1 / 1e6))
print(' + RefineDet Number of params: %.2fM' % (total2 / 1e6))
print_model_parm_flops(net1, 300)
print_model_parm_flops(net2, 320)
def train():
best_map = 0
if args.dataset == 'VOC':
if args.dataset_root == COCO_ROOT:
parser.error('Must specify dataset if specifying dataset_root')
cfg = voc_refinedet[args.input_size]
dataset = VOCDetection(root=args.dataset_root,
transform=SSDAugmentation(
cfg['min_dim'], MEANS))
refinedet_net = build_refinedet('train', cfg['min_dim'],
cfg['num_classes'])
net = refinedet_net
# print(net)
#input()
if args.cuda:
# net = torch.nn.DataParallel(refinedet_net)
net = net.to(device)
cudnn.benchmark = True
if args.resume:
print('Resuming training, loading {}...'.format(args.resume))
refinedet_net.load_weights(args.resume)
else:
#vgg_weights = torch.load(args.save_folder + args.basenet)
vgg_weights = torch.load(args.basenet)
print('Loading base network...')
refinedet_net.vgg.load_state_dict(vgg_weights)
# if args.cuda:
# net = net.cuda()
if not args.resume:
print('Initializing weights...')
# initialize newly added layers' weights with xavier method
refinedet_net.extras.apply(weights_init)
refinedet_net.arm_loc.apply(weights_init)
refinedet_net.arm_conf.apply(weights_init)
refinedet_net.odm_loc.apply(weights_init)
refinedet_net.odm_conf.apply(weights_init)
#refinedet_net.tcb.apply(weights_init)
refinedet_net.tcb0.apply(weights_init)
refinedet_net.tcb1.apply(weights_init)
refinedet_net.tcb2.apply(weights_init)
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
arm_criterion = RefineDetMultiBoxLoss(2, 0.5, True, 0, True, 3, 0.5, False,
args.cuda)
odm_criterion = RefineDetMultiBoxLoss(cfg['num_classes'],
0.5,
True,
0,
True,
3,
0.5,
False,
args.cuda,
use_ARM=True)
# net.train()
# loss counters
arm_loc_loss = 0
arm_conf_loss = 0
odm_loc_loss = 0
odm_conf_loss = 0
epoch = 0
print('Loading the dataset...')
epoch_size = len(dataset) // args.batch_size
print('Training RefineDet on:', dataset.name)
print('Using the specified args:')
print(args)
step_index = 0
data_loader = data.DataLoader(dataset,
args.batch_size,
num_workers=args.num_workers,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
# create batch iterator
iteration = 0
for epoch in range(args.epochs):
net.phase = 'train'
net.train()
for batch_i, (images, targets) in enumerate(data_loader):
iteration += 1
if iteration in cfg['lr_steps']:
step_index += 1
adjust_learning_rate(optimizer, args.gamma, step_index)
# load train data
images = images.to(device)
targets = [ann.to(device) for ann in targets]
# forward
t0 = time.time()
out = net(images)
# backprop
optimizer.zero_grad()
arm_loss_l, arm_loss_c = arm_criterion(out, targets)
odm_loss_l, odm_loss_c = odm_criterion(out, targets)
#input()
arm_loss = arm_loss_l + arm_loss_c
odm_loss = odm_loss_l + odm_loss_c
loss = arm_loss + odm_loss
loss.backward()
optimizer.step()
t1 = time.time()
arm_loc_loss += arm_loss_l.item()
arm_conf_loss += arm_loss_c.item()
odm_loc_loss += odm_loss_l.item()
odm_conf_loss += odm_loss_c.item()
if iteration % 10 == 0:
print('timer: %.4f sec.' % (t1 - t0))
print('epoch '+repr(epoch+1)+': iter ' + repr(iteration) + ' || ARM_L Loss: %.4f ARM_C Loss: %.4f ODM_L Loss: %.4f ODM_C Loss: %.4f ||' \
% (arm_loss_l.item(), arm_loss_c.item(), odm_loss_l.item(), odm_loss_c.item()), end=' ')
if epoch % args.evaluation_interval == 0:
print("\n---- Evaluating Model ----")
map = evaluate(model=net,
save_folder=args.save_folders,
cuda=args.cuda,
top_k=5,
im_size=320,
thresh=0.05,
dataset_mean=((104, 117, 123)))
if map > best_map: #取最好的map保存
torch.save(net.state_dict(),
f"checkpoints/RefineDet320_%d.pth" % (epoch + 1))
best_map = map