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