def test(): # get device device = get_device(0) # load net num_classes = len(VOC_CLASSES) testset = VOCDetection(args.voc_root, [('2007', 'test')], None, VOCAnnotationTransform()) mean = config.MEANS cfg = config.voc_ab if args.version == 'yolo_v2': net = myYOLOv2(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=config.ANCHOR_SIZE) print('Let us test yolo-v2 on the VOC0712 dataset ......') elif args.version == 'yolo_v3': from models.yolo_v3 import myYOLOv3 net = myYOLOv3(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=config.MULTI_ANCHOR_SIZE) 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), thresh=args.visual_threshold)
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 = len(VOC_CLASSES) testset = VOCDetection(args.voc_root, [('2007', 'test')], None, VOCAnnotationTransform()) cfg = config.voc_ab 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) print('Let us test yolo-v2 on the VOC0712 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) elif args.version == 'slim_yolo_v2': from models.slim_yolo_v2 import SlimYOLOv2 net = SlimYOLOv2(device, input_size=cfg['min_dim'], num_classes=num_classes, anchor_size=config.ANCHOR_SIZE) print('Let us test slim-yolo-v2 on the VOC0712 dataset ......') 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.TINY_MULTI_ANCHOR_SIZE) print('Let us test tiny-yolo-v3 on the VOC0712 dataset ......') net.load_state_dict(torch.load(args.trained_model, map_location=device)) 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)
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 == 'slim_yolo_v2': from models.slim_yolo_v2 import SlimYOLOv2 net = SlimYOLOv2(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.TINY_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)
def run(): args = parse_args() # use cuda if args.cuda: device = torch.device("cuda") else: device = torch.device("cpu") input_size = [args.input_size, args.input_size] # load net if args.version == 'yolo_v2': from models.yolo_v2 import myYOLOv2 anchor_size = config.ANCHOR_SIZE_COCO net = myYOLOv2(device, input_size=input_size, num_classes=80, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh, anchor_size=anchor_size) elif args.version == 'yolo_v3': from models.yolo_v3 import myYOLOv3 anchor_size = config.MULTI_ANCHOR_SIZE_COCO net = myYOLOv3(device, input_size=input_size, num_classes=80, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh, anchor_size=anchor_size) elif args.version == 'slim_yolo_v2': from models.slim_yolo_v2 import SlimYOLOv2 anchor_size = config.ANCHOR_SIZE_COCO net = SlimYOLOv2(device, input_size=input_size, num_classes=80, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh, anchor_size=anchor_size) elif args.version == 'tiny_yolo_v3': from models.tiny_yolo_v3 import YOLOv3tiny anchor_size = config.TINY_MULTI_ANCHOR_SIZE_COCO net = YOLOv3tiny(device, input_size=input_size, num_classes=80, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh, anchor_size=anchor_size) net.load_state_dict(torch.load(args.trained_model, map_location='cuda')) net.to(device).eval() print('Finished loading model!') # run if args.mode == 'camera': detect(net, device, BaseTransform(net.input_size), thresh=args.visual_threshold, mode=args.mode) elif args.mode == 'image': detect(net, device, BaseTransform(net.input_size), thresh=args.visual_threshold, mode=args.mode, path_to_img=args.path_to_img) elif args.mode == 'video': detect(net, device, BaseTransform(net.input_size), thresh=args.visual_threshold, mode=args.mode, path_to_vid=args.path_to_vid, path_to_save=args.path_to_saveVid)
def test(): # get device device = get_device(0) # load net num_classes = 80 anchor_size = config.ANCHOR_SIZE_COCO 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) mean = config.MEANS elif args.dataset == 'VOC': cfg = config.voc_ab testset = VOCDetection(VOC_ROOT, [('2007', 'test')], None, VOCAnnotationTransform()) mean = config.MEANS if args.version == 'yolo_v2': from models.yolo_v2 import myYOLOv2 net = myYOLOv2(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=anchor_size) 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, trainable=False, anchor_size=anchor_size) 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, trainable=False, anchor_size=config.ANCHOR_SIZE) 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, trainable=False, anchor_size=config.MULTI_ANCHOR_SIZE) 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), thresh=args.visual_threshold)
write_voc_results_file(box_list, dataset) do_python_eval(output_dir) if __name__ == '__main__': num_classes = len(labelmap) device = get_device(args.gpu_ind) cfg = config.voc_ab if args.version == 'yolo_v2': from models.yolo_v2 import myYOLOv2 net = myYOLOv2(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=config.ANCHOR_SIZE) elif args.version == 'yolo_v3': from models.yolo_v3 import myYOLOv3 net = myYOLOv3(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=config.MULTI_ANCHOR_SIZE) 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, trainable=False, anchor_size=config.ANCHOR_SIZE) print('Let us eval tiny-yolo-v2 on the VOC0712 dataset ......') 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, trainable=False, anchor_size=config.MULTI_ANCHOR_SIZE) print('Let us eval tiny-yolo-v3 on the VOC0712 dataset ......') # load net net.load_state_dict(torch.load(args.trained_model, map_location='cuda')) net.eval() print('Finished loading model!')
device = torch.device("cuda") else: device = torch.device("cpu") if args.version == 'yolo_v2': from models.yolo_v2 import myYOLOv2 model = myYOLOv2(device, input_size=cfg['min_dim'], num_classes=args.num_classes, anchor_size=ANCHOR_SIZE_COCO) print('Let us test yolo-v2 on the MSCOCO dataset ......') elif args.version == 'yolo_v3': from models.yolo_v3 import myYOLOv3 model = myYOLOv3(device, input_size=cfg['min_dim'], num_classes=args.num_classes, anchor_size=MULTI_ANCHOR_SIZE_COCO) elif args.version == 'tiny_yolo_v2': from models.tiny_yolo_v2 import YOLOv2tiny model = YOLOv2tiny(device, input_size=cfg['min_dim'], num_classes=args.num_classes, anchor_size=ANCHOR_SIZE_COCO) elif args.version == 'tiny_yolo_v3': from models.tiny_yolo_v3 import YOLOv3tiny model = YOLOv3tiny(device, input_size=cfg['min_dim'], num_classes=args.num_classes,
if args.high_resolution == 1: hr = True cfg = coco_ab if args.version == 'yolo_v2': from models.yolo_v2 import myYOLOv2 total_anchor_size = tools.get_total_anchor_size(name='COCO') yolo_net = myYOLOv2(device, input_size=cfg['min_dim'], num_classes=args.num_classes, trainable=True, anchor_size=total_anchor_size, hr=hr) print('Let us train yolo-v2 on the MSCOCO dataset ......') elif args.version == 'yolo_v3': from models.yolo_v3 import myYOLOv3 total_anchor_size = tools.get_total_anchor_size(multi_scale=True, name='COCO') yolo_net = myYOLOv3(device, input_size=cfg['min_dim'], num_classes=args.num_classes, trainable=True, anchor_size=total_anchor_size, hr=hr) print('Let us train yolo-v3 on the MSCOCO dataset ......') train(yolo_net, device)
def train(): args = parse_args().parse_args(args=[]) path_to_save = os.path.join(args.save_folder, args.dataset, args.version) os.makedirs(path_to_save, exist_ok=True) # use hi-res backbone if args.high_resolution: print('use hi-res backbone') hr = True else: hr = False # cuda if args.cuda: print('use cuda') cudnn.benchmark = True device = torch.device("cuda", 0) else: device = torch.device("cpu") # multi-scale if args.multi_scale: print('use the multi-scale trick ...') train_size = config_0['NIH_pancreas_data_aimshape'] val_size = config_0['NIH_pancreas_data_aimshape'] else: train_size = config_0['NIH_pancreas_data_aimshape'] val_size = config_0['NIH_pancreas_data_aimshape'] cfg = train_cfg # dataset and evaluator print("Setting Arguments.. : ", args) print("----------------------------------------------------------") print('Loading the dataset...') if args.dataset == 'pnens': num_classes = 1 # dataset_pkl = load_from_pkl(r'/data/liyi219/pnens_3D_data/after_dealing/pre_order0_128_128_64_new.pkl') dataset_pkl = load_from_pkl( r'E:\ly\pnens_data\nas_data\v1_data\NIH\pre_order0_128_128_64_new.pkl' ) dataset = datanih(dataset_pkl) evaluator = myEvaluator(dataset=dataset, data_root="/data/data4T/ly/data/pnens_3D", img_size=val_size, device=device, transform=BaseTransform(val_size), labelmap=('pnens')) else: print('unknow dataset !! Only support voc and coco !!') exit(0) print('Training model on: yolov3_3D') print('The dataset size:', len(dataset)) print("----------------------------------------------------------") # dataloader dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=False, collate_fn=detection_collate, num_workers=args.num_workers, pin_memory=True) if args.version == 'yolo_v3': from models.yolo_v3 import myYOLOv3 # anchor_size = MULTI_ANCHOR_SIZE if args.dataset == 'voc' else MULTI_ANCHOR_SIZE_COCO anchor_size = anchor_size_3D_try yolo_net = myYOLOv3(device, input_size=train_size, num_classes=num_classes, trainable=True, anchor_size=anchor_size, hr=hr) print('Let us train yolo_v3 on the %s dataset ......' % (args.dataset)) else: print('Unknown version !!!') exit() model = yolo_net model.to(device).train() # use tfboard if args.tfboard: print('use tensorboard') from torch.utils.tensorboard import SummaryWriter c_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())) log_path = os.path.join('log/coco/', args.version, c_time) os.makedirs(log_path, exist_ok=True) writer = SummaryWriter(log_path) # keep training if args.resume is not None: print('keep training model: %s' % (args.resume)) model.load_state_dict(torch.load(args.resume, map_location=device)) # optimizer setup base_lr = args.lr tmp_lr = base_lr optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) max_epoch = cfg['max_epoch'] epoch_size = len(dataset) // args.batch_size # start training loop t0 = time.time() for epoch in range(args.start_epoch, max_epoch): # if epoch == 1: break # use cos lr if args.cos and epoch > 20 and epoch <= max_epoch - 20: # use cos lr tmp_lr = 0.00001 + 0.5 * (base_lr - 0.00001) * ( 1 + math.cos(math.pi * (epoch - 20) * 1. / (max_epoch - 20))) set_lr(optimizer, tmp_lr) elif args.cos and epoch > max_epoch - 20: tmp_lr = 0.00001 set_lr(optimizer, tmp_lr) # use step lr else: if epoch in cfg['lr_epoch']: tmp_lr = tmp_lr * 0.1 set_lr(optimizer, tmp_lr) for iter_i, (images, targets) in enumerate(dataloader): """ images [1, 1, 128, 128, 64] targets [y1, y2, x1, x2, z1, z2, 0] 相对 """ # targets [x1, x2, y1, y2, z1, z2, 有无物体(0 or 1)] 相对坐标[0~1] # if iter_i == 0: break if not args.no_warm_up: if epoch < args.wp_epoch: tmp_lr = base_lr * pow((iter_i + epoch * epoch_size) * 1. / (args.wp_epoch * epoch_size), 4) # tmp_lr = 1e-6 + (base_lr-1e-6) * (iter_i+epoch*epoch_size) / (epoch_size * (args.wp_epoch)) set_lr(optimizer, tmp_lr) elif epoch == args.wp_epoch and iter_i == 0: tmp_lr = base_lr set_lr(optimizer, tmp_lr) # to device images = images.to(device) # multi-scale trick if iter_i % 10 == 0 and iter_i > 0 and args.multi_scale: # randomly choose a new size size = random.randint(10, 19) * 32 train_size = [size, size] model.set_grid(train_size) if args.multi_scale: # interpolate # 上采样 images = torch.nn.functional.interpolate(images, size=train_size, mode='bilinear', align_corners=False) # make labels # print(targets) targets = [label.tolist() for label in targets] targets = tools.multi_gt_creator3D(input_size=train_size, strides=model.stride, label_lists=targets, anchor_size=anchor_size) # [batch_index, featuremap_index, grid_x* grid_y * grid_z * ab_ind, 参数(15)] # (15): [obj, class, tx, ty, tz, th, tw, td, weight, (xmin, ymin, zmin, ymax, xmax, zmax)(绝对)] targets = torch.tensor(targets).float().to(device) # forward and loss conf_loss, cls_loss, txtytwth_loss, total_loss, dice_loss = model( images, target=targets) # backprop total_loss.backward() optimizer.step() optimizer.zero_grad() # display if iter_i % 10 == 0: if args.tfboard: # viz loss writer.add_scalar('object loss', conf_loss.item(), iter_i + epoch * epoch_size) writer.add_scalar('class loss', cls_loss.item(), iter_i + epoch * epoch_size) writer.add_scalar('local loss', txtytwth_loss.item(), iter_i + epoch * epoch_size) t1 = time.time() print( '[Epoch %d/%d][Iter %d/%d][lr %.6f]' '[Loss: obj %.2f ||dice_loss %.2f || cls %.2f || bbox %.2f || total %.2f || size %d || time: %.2f]' % (epoch + 1, max_epoch, iter_i, epoch_size, tmp_lr, conf_loss.item(), dice_loss.item(), cls_loss.item(), txtytwth_loss.item(), total_loss.item(), train_size[0], t1 - t0), flush=True) t0 = time.time() # evaluation if (epoch + 1) % args.eval_epoch == 0: model.trainable = False model.set_grid(val_size) model.eval() # evaluate # evaluator.evaluate(model) # convert to training mode. model.trainable = True model.set_grid(train_size) model.train() # save model if (epoch + 1) % 5 == 0: print('Saving state, saving in {} epoch{}:'.format( os.path.join(path_to_save, args.version), epoch + 1, )) torch.save( model.state_dict(), os.path.join(path_to_save, args.version + '_' + repr(epoch + 1) + '.pth'))
else: device = torch.device("cpu") # input size input_size = [args.input_size, args.input_size] # load net if args.version == 'yolo_v2': from models.yolo_v2 import myYOLOv2 anchor_size = ANCHOR_SIZE if args.dataset == 'voc' else ANCHOR_SIZE_COCO net = myYOLOv2(device, input_size=input_size, num_classes=num_classes, anchor_size=anchor_size) elif args.version == 'yolo_v3': from models.yolo_v3 import myYOLOv3 anchor_size = MULTI_ANCHOR_SIZE if args.dataset == 'voc' else MULTI_ANCHOR_SIZE_COCO net = myYOLOv3(device, input_size=input_size, num_classes=num_classes, anchor_size=anchor_size) elif args.version == 'yolo_v3_spp': from models.yolo_v3_spp import myYOLOv3Spp anchor_size = MULTI_ANCHOR_SIZE if args.dataset == 'voc' else MULTI_ANCHOR_SIZE_COCO net = myYOLOv3Spp(device, input_size=input_size, num_classes=num_classes, anchor_size=anchor_size) elif args.version == 'slim_yolo_v2': from models.slim_yolo_v2 import SlimYOLOv2 anchor_size = ANCHOR_SIZE if args.dataset == 'voc' else ANCHOR_SIZE_COCO net = SlimYOLOv2(device, input_size=input_size, num_classes=num_classes, anchor_size=anchor_size) elif args.version == 'tiny_yolo_v3': from models.tiny_yolo_v3 import YOLOv3tiny anchor_size = TINY_MULTI_ANCHOR_SIZE if args.dataset == 'voc' else TINY_MULTI_ANCHOR_SIZE_COCO net = YOLOv3tiny(device, input_size=input_size, num_classes=num_classes, anchor_size=anchor_size)
def train(): args = parse_args() path_to_save = os.path.join(args.save_folder, args.version) os.makedirs(path_to_save, exist_ok=True) hr = False if args.high_resolution: print('use hi-res backbone') hr = True cfg = voc_ab if args.cuda: print('use cuda') cudnn.benchmark = True device = torch.device("cuda") else: device = torch.device("cpu") # use multi-scale trick if args.multi_scale: print('use multi-scale trick.') input_size = [608, 608] dataset = VOCDetection(root=args.dataset_root, transform=SSDAugmentation([608, 608], mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229))) else: input_size = cfg['min_dim'] dataset = VOCDetection(root=args.dataset_root, transform=SSDAugmentation(cfg['min_dim'], mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229))) # build model if args.version == 'yolo_v2': from models.yolo_v2 import myYOLOv2 total_anchor_size = tools.get_total_anchor_size() yolo_net = myYOLOv2(device, input_size=input_size, num_classes=args.num_classes, trainable=True, anchor_size=total_anchor_size, hr=hr) print('Let us train yolo-v2 on the VOC0712 dataset ......') elif args.version == 'yolo_v3': from models.yolo_v3 import myYOLOv3 total_anchor_size = tools.get_total_anchor_size(multi_level=True, version='yolo_v3') yolo_net = myYOLOv3(device, input_size=input_size, num_classes=args.num_classes, trainable=True, anchor_size=total_anchor_size, hr=hr) print('Let us train yolo-v3 on the VOC0712 dataset ......') elif args.version == 'slim_yolo_v2': from models.slim_yolo_v2 import SlimYOLOv2 total_anchor_size = tools.get_total_anchor_size() yolo_net = SlimYOLOv2(device, input_size=input_size, num_classes=args.num_classes, trainable=True, anchor_size=total_anchor_size, hr=hr) print('Let us train slim-yolo-v2 on the VOC0712 dataset ......') elif args.version == 'tiny_yolo_v3': from models.tiny_yolo_v3 import YOLOv3tiny total_anchor_size = tools.get_total_anchor_size(multi_level=True, version='tiny_yolo_v3') yolo_net = YOLOv3tiny(device, input_size=input_size, num_classes=args.num_classes, trainable=True, anchor_size=total_anchor_size, hr=hr) print('Let us train tiny-yolo-v3 on the VOC0712 dataset ......') else: print('Unknown version !!!') exit() # finetune the model trained on COCO if args.resume is not None: print('finetune COCO trained ') yolo_net.load_state_dict(torch.load(args.resume, map_location=device), strict=False) # use tfboard if args.tfboard: print('use tensorboard') from torch.utils.tensorboard import SummaryWriter c_time = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())) log_path = os.path.join('log/voc/', args.version, c_time) os.makedirs(log_path, exist_ok=True) writer = SummaryWriter(log_path) print("----------------------------------------Object Detection--------------------------------------------") model = yolo_net model.to(device) base_lr = args.lr tmp_lr = base_lr optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # loss counters print("----------------------------------------------------------") print("Let's train OD network !") print('Training on:', dataset.name) print('The dataset size:', len(dataset)) print("----------------------------------------------------------") epoch_size = len(dataset) // args.batch_size max_epoch = cfg['max_epoch'] 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 t0 = time.time() # start training for epoch in range(max_epoch): # use cos lr if args.cos and epoch > 20 and epoch <= max_epoch - 20: # use cos lr tmp_lr = 0.00001 + 0.5*(base_lr-0.00001)*(1+math.cos(math.pi*(epoch-20)*1./ (max_epoch-20))) set_lr(optimizer, tmp_lr) elif args.cos and epoch > max_epoch - 20: tmp_lr = 0.00001 set_lr(optimizer, tmp_lr) # use step lr else: if epoch in cfg['lr_epoch']: tmp_lr = tmp_lr * 0.1 set_lr(optimizer, tmp_lr) for iter_i, (images, targets) in enumerate(data_loader): # WarmUp strategy for learning rate if not args.no_warm_up: if epoch < args.wp_epoch: tmp_lr = base_lr * pow((iter_i+epoch*epoch_size)*1. / (args.wp_epoch*epoch_size), 4) # tmp_lr = 1e-6 + (base_lr-1e-6) * (iter_i+epoch*epoch_size) / (epoch_size * (args.wp_epoch)) set_lr(optimizer, tmp_lr) elif epoch == args.wp_epoch and iter_i == 0: tmp_lr = base_lr set_lr(optimizer, tmp_lr) targets = [label.tolist() for label in targets] # make train label if args.version == 'yolo_v2' or args.version == 'slim_yolo_v2': targets = tools.gt_creator(input_size, yolo_net.stride, targets, version=args.version) elif args.version == 'yolo_v3' or args.version == 'tiny_yolo_v3': targets = tools.multi_gt_creator(input_size, yolo_net.stride, targets, version=args.version) # to device images = images.to(device) targets = torch.tensor(targets).float().to(device) # forward and loss conf_loss, cls_loss, txtytwth_loss, total_loss = model(images, target=targets) # backprop and update total_loss.backward() optimizer.step() optimizer.zero_grad() if iter_i % 10 == 0: if args.tfboard: # viz loss writer.add_scalar('object loss', conf_loss.item(), iter_i + epoch * epoch_size) writer.add_scalar('class loss', cls_loss.item(), iter_i + epoch * epoch_size) writer.add_scalar('local loss', txtytwth_loss.item(), iter_i + epoch * epoch_size) t1 = time.time() print('[Epoch %d/%d][Iter %d/%d][lr %.6f]' '[Loss: obj %.2f || cls %.2f || bbox %.2f || total %.2f || size %d || time: %.2f]' % (epoch+1, max_epoch, iter_i, epoch_size, tmp_lr, conf_loss.item(), cls_loss.item(), txtytwth_loss.item(), total_loss.item(), input_size[0], t1-t0), flush=True) t0 = time.time() # multi-scale trick if iter_i % 10 == 0 and iter_i > 0 and args.multi_scale: size = random.randint(10, 19) * 32 input_size = [size, size] model.set_grid(input_size) # change input dim # But this operation will report bugs when we use more workers in data loader, so I have to use 0 workers. # I don't know how to make it suit more workers, and I'm trying to solve this question. data_loader.dataset.reset_transform(SSDAugmentation(input_size, mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229))) if (epoch + 1) % 10 == 0: print('Saving state, epoch:', epoch + 1) torch.save(model.state_dict(), os.path.join(path_to_save, args.version + '_' + repr(epoch + 1) + '.pth') )
def run(): args = parse_args() if args.cuda: device = torch.device("cuda") else: device = torch.device("cpu") if args.setup == 'VOC': print('use VOC style') cfg = config.voc_ab num_classes = 20 elif args.setup == 'COCO': print('use COCO style') cfg = config.coco_ab num_classes = 80 else: print('Only support VOC and COCO !!!') exit(0) if args.version == 'yolo_v2': from models.yolo_v2 import myYOLOv2 if args.setup == 'VOC': anchor_size = config.ANCHOR_SIZE else: anchor_size = config.ANCHOR_SIZE_COCO net = myYOLOv2(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=anchor_size) elif args.version == 'yolo_v3': from models.yolo_v3 import myYOLOv3 if args.setup == 'VOC': anchor_size = config.MULTI_ANCHOR_SIZE else: anchor_size = config.MULTI_ANCHOR_SIZE_COCO net = myYOLOv3(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=anchor_size) elif args.version == 'tiny_yolo_v2': from models.tiny_yolo_v2 import YOLOv2tiny if args.setup == 'VOC': anchor_size = config.ANCHOR_SIZE else: anchor_size = config.ANCHOR_SIZE_COCO net = YOLOv2tiny(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=anchor_size) elif args.version == 'tiny_yolo_v3': from models.tiny_yolo_v3 import YOLOv3tiny if args.setup == 'VOC': anchor_size = config.MULTI_ANCHOR_SIZE else: anchor_size = config.MULTI_ANCHOR_SIZE_COCO net = YOLOv3tiny(device, input_size=cfg['min_dim'], num_classes=num_classes, trainable=False, anchor_size=anchor_size) net.load_state_dict(torch.load(args.trained_model, map_location='cuda')) net.to(device).eval() print('Finished loading model!') # run if args.mode == 'image': detect(net, device, BaseTransform(net.input_size, mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229)), mode=args.mode, path_to_img=args.path_to_img, setup=args.setup) elif args.mode == 'video': detect(net, device, BaseTransform(net.input_size, mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229)), mode=args.mode, path_to_vid=args.path_to_vid, path_to_save=args.path_to_saveVid, setup=args.setup)
def train(): args = parse_args() path_to_save = os.path.join(args.save_folder, args.dataset, args.version) os.makedirs(path_to_save, exist_ok=True) # use hi-res backbone if args.high_resolution: print('use hi-res backbone') hr = True else: hr = False # cuda if args.cuda: print('use cuda') cudnn.benchmark = True device = torch.device("cuda") else: device = torch.device("cpu") # multi-scale if args.multi_scale: print('use the multi-scale trick ...') train_size = [640, 640] val_size = [416, 416] else: train_size = [416, 416] val_size = [416, 416] cfg = train_cfg # dataset and evaluator print("Setting Arguments.. : ", args) print("----------------------------------------------------------") print('Loading the dataset...') if args.dataset == 'voc': data_dir = VOC_ROOT num_classes = 20 dataset = VOCDetection(root=data_dir, transform=SSDAugmentation(train_size)) evaluator = VOCAPIEvaluator(data_root=data_dir, img_size=val_size, device=device, transform=BaseTransform(val_size), labelmap=VOC_CLASSES) elif args.dataset == 'coco': data_dir = coco_root num_classes = 80 dataset = COCODataset(data_dir=data_dir, img_size=train_size[0], transform=SSDAugmentation(train_size), debug=args.debug) evaluator = COCOAPIEvaluator(data_dir=data_dir, img_size=val_size, device=device, transform=BaseTransform(val_size)) else: print('unknow dataset !! Only support voc and coco !!') exit(0) print('Training model on:', dataset.name) print('The dataset size:', len(dataset)) print("----------------------------------------------------------") # dataloader dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True, collate_fn=detection_collate, num_workers=args.num_workers, pin_memory=True) # build model if args.version == 'yolo_v2': from models.yolo_v2 import myYOLOv2 anchor_size = ANCHOR_SIZE if args.dataset == 'voc' else ANCHOR_SIZE_COCO yolo_net = myYOLOv2(device, input_size=train_size, num_classes=num_classes, trainable=True, anchor_size=anchor_size, hr=hr) print('Let us train yolo_v2 on the %s dataset ......' % (args.dataset)) elif args.version == 'yolo_v3': from models.yolo_v3 import myYOLOv3 anchor_size = MULTI_ANCHOR_SIZE if args.dataset == 'voc' else MULTI_ANCHOR_SIZE_COCO yolo_net = myYOLOv3(device, input_size=train_size, num_classes=num_classes, trainable=True, anchor_size=anchor_size, hr=hr) print('Let us train yolo_v3 on the %s dataset ......' % (args.dataset)) elif args.version == 'yolo_v3_spp': from models.yolo_v3_spp import myYOLOv3Spp anchor_size = MULTI_ANCHOR_SIZE if args.dataset == 'voc' else MULTI_ANCHOR_SIZE_COCO yolo_net = myYOLOv3Spp(device, input_size=train_size, num_classes=num_classes, trainable=True, anchor_size=anchor_size, hr=hr) print('Let us train yolo_v3_spp on the %s dataset ......' % (args.dataset)) elif args.version == 'slim_yolo_v2': from models.slim_yolo_v2 import SlimYOLOv2 anchor_size = ANCHOR_SIZE if args.dataset == 'voc' else ANCHOR_SIZE_COCO yolo_net = SlimYOLOv2(device, input_size=train_size, num_classes=num_classes, trainable=True, anchor_size=anchor_size, hr=hr) print('Let us train slim_yolo_v2 on the %s dataset ......' % (args.dataset)) elif args.version == 'tiny_yolo_v3': from models.tiny_yolo_v3 import YOLOv3tiny anchor_size = TINY_MULTI_ANCHOR_SIZE if args.dataset == 'voc' else TINY_MULTI_ANCHOR_SIZE_COCO yolo_net = YOLOv3tiny(device, input_size=train_size, num_classes=num_classes, trainable=True, anchor_size=anchor_size, hr=hr) print('Let us train tiny_yolo_v3 on the %s dataset ......' % (args.dataset)) else: print('Unknown version !!!') exit() model = yolo_net model.to(device).train() # use tfboard if args.tfboard: print('use tensorboard') from torch.utils.tensorboard import SummaryWriter c_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())) log_path = os.path.join('log/coco/', args.version, c_time) os.makedirs(log_path, exist_ok=True) writer = SummaryWriter(log_path) # keep training if args.resume is not None: print('keep training model: %s' % (args.resume)) model.load_state_dict(torch.load(args.resume, map_location=device)) # optimizer setup base_lr = args.lr tmp_lr = base_lr optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) max_epoch = cfg['max_epoch'] epoch_size = len(dataset) // args.batch_size # start training loop t0 = time.time() for epoch in range(args.start_epoch, max_epoch): # use cos lr if args.cos and epoch > 20 and epoch <= max_epoch - 20: # use cos lr tmp_lr = 0.00001 + 0.5 * (base_lr - 0.00001) * ( 1 + math.cos(math.pi * (epoch - 20) * 1. / (max_epoch - 20))) set_lr(optimizer, tmp_lr) elif args.cos and epoch > max_epoch - 20: tmp_lr = 0.00001 set_lr(optimizer, tmp_lr) # use step lr else: if epoch in cfg['lr_epoch']: tmp_lr = tmp_lr * 0.1 set_lr(optimizer, tmp_lr) for iter_i, (images, targets) in enumerate(dataloader): # WarmUp strategy for learning rate if not args.no_warm_up: if epoch < args.wp_epoch: tmp_lr = base_lr * pow((iter_i + epoch * epoch_size) * 1. / (args.wp_epoch * epoch_size), 4) # tmp_lr = 1e-6 + (base_lr-1e-6) * (iter_i+epoch*epoch_size) / (epoch_size * (args.wp_epoch)) set_lr(optimizer, tmp_lr) elif epoch == args.wp_epoch and iter_i == 0: tmp_lr = base_lr set_lr(optimizer, tmp_lr) # to device images = images.to(device) # multi-scale trick if iter_i % 10 == 0 and iter_i > 0 and args.multi_scale: # randomly choose a new size size = random.randint(10, 19) * 32 train_size = [size, size] model.set_grid(train_size) if args.multi_scale: # interpolate images = torch.nn.functional.interpolate(images, size=train_size, mode='bilinear', align_corners=False) # make labels targets = [label.tolist() for label in targets] if args.version == 'yolo_v2' or args.version == 'slim_yolo_v2': targets = tools.gt_creator(input_size=train_size, stride=yolo_net.stride, label_lists=targets, anchor_size=anchor_size) else: targets = tools.multi_gt_creator(input_size=train_size, strides=yolo_net.stride, label_lists=targets, anchor_size=anchor_size) targets = torch.tensor(targets).float().to(device) # forward and loss conf_loss, cls_loss, txtytwth_loss, total_loss = model( images, target=targets) # backprop total_loss.backward() optimizer.step() optimizer.zero_grad() # display if iter_i % 10 == 0: if args.tfboard: # viz loss writer.add_scalar('object loss', conf_loss.item(), iter_i + epoch * epoch_size) writer.add_scalar('class loss', cls_loss.item(), iter_i + epoch * epoch_size) writer.add_scalar('local loss', txtytwth_loss.item(), iter_i + epoch * epoch_size) t1 = time.time() print( '[Epoch %d/%d][Iter %d/%d][lr %.6f]' '[Loss: obj %.2f || cls %.2f || bbox %.2f || total %.2f || size %d || time: %.2f]' % (epoch + 1, max_epoch, iter_i, epoch_size, tmp_lr, conf_loss.item(), cls_loss.item(), txtytwth_loss.item(), total_loss.item(), train_size[0], t1 - t0), flush=True) t0 = time.time() # evaluation if (epoch + 1) % args.eval_epoch == 0: model.trainable = False model.set_grid(val_size) model.eval() # evaluate evaluator.evaluate(model) # convert to training mode. model.trainable = True model.set_grid(train_size) model.train() # save model if (epoch + 1) % 10 == 0: print('Saving state, epoch:', epoch + 1) torch.save( model.state_dict(), os.path.join(path_to_save, args.version + '_' + repr(epoch + 1) + '.pth'))
img_size=input_size[0]) elif args.dataset == 'pnens': num_classes = 1 class_names = ('pnens') # dataset2 = load_from_pkl(r'/data/liyi219/pnens_3D_data/after_dealing/pre_order0_128_128_64_new.pkl') dataset2 = load_from_pkl('C:\\Users\\hazy\\Desktop\\fsdownload\\\pre_order0_128_128_64_new.pkl') # dataset = rechange(dataset2) dataset = datanih(dataset2) class_colors = [(np.random.randint(255), np.random.randint(255), np.random.randint(255)) for _ in range(num_classes)] # load net from models.yolo_v3 import myYOLOv3 anchor_size = anchor_size_3D_try net = myYOLOv3(device, input_size=input_size, num_classes=1, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh, anchor_size=anchor_size) # net.load_state_dict(torch.load(args.trained_model, map_location=device)) net.to(device).eval() print('Finished loading model!') # evaluation test(net=net, device=device, testset=dataset, transform=BaseTransform(input_size), thresh=args.visual_threshold, class_colors=class_colors, class_names=class_names, class_indexs=[0], dataset_data_kind=args.dataset
def see_bbox(): args = parse_args().parse_args(args=[]) device = torch.device("cpu") input_size = [128, 128, 64] # dataset2 = load_from_pkl(r'/data/liyi219/pnens_3D_data/after_dealing/pre_order0_128_128_64_new.pkl') dataset2 = load_from_pkl( r'E:\ly\pnens_data\nas_data\v1_data\NIH\pre_order0_128_128_64_new.pkl') dataset = datanih(dataset2) # load net from models.yolo_v3 import myYOLOv3 anchor_size = anchor_size_3D_try net = myYOLOv3(device, input_size=input_size, num_classes=1, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh, anchor_size=anchor_size) net.load_state_dict( torch.load(os.path.join(args.trained_model + 'yolo_v3_200.pth'), map_location=device)) net.to(device).eval() num_images = len(dataset) result = [] images = [] masks = [] bbox_gt = [] for index in range(num_images): if index == 1: break # target im, gt = dataset.__getitem__(index) targets = get_targets(gt) # gt [[obj, class, tx, ty, tz, tw, th, td, weight, (xmin, ymin, zmin, xmax, ymax, zmax)(绝对)]] gt = targets[np.where(targets == 1)[0:2]] txtytztwthtd = gt[0, 2:8] gt_obj = targets[:, :, 0] map_obj = obj_map_get(gt_obj) print('Testing image {:d}/{:d}....'.format(index + 1, num_images)) img, _, mask, height, width, depth = dataset.pull_item(index) bb = dataset.get_bbox_juedui(index) # [y1, y2, x1, x2, z1, z2] 绝对 bbox_gt.append(bb) scale = np.array([[width, height, depth, width, height, depth]]) images.append(img) masks.append(mask) # to tensor x = torch.tensor(img) x = x.unsqueeze(0).to(device) x = x.unsqueeze(0).to(device) t0 = time.time() # forward bboxes, scores, _ = net(x) # bboxes[x1, y1, z1, x2, y2, z2] 相对 bboxes = bboxes * scale # to绝对 if not scores.size: print('{} has not bbox'.format(index)) continue best_scores = np.argmax(scores, axis=0) bboxes = bboxes[best_scores] # print(scores[best_scores]) # print([bboxes, best_scores]) result.append([bboxes, best_scores]) print("detection {} time used ".format(index), time.time() - t0, "s") if False: pred_1 = torch.sigmoid(pred_1) pred_1 = pred_1.data.cpu().numpy()[0, 0] img3D = pred_1 vol = mlab.pipeline.volume(mlab.pipeline.scalar_field(img3D), name='3-d ultrasound ') mlab.colorbar(orientation='vertical') mlab.show() show3Dslice(img3D) img_resize = resize3D(img3D, [128, 64, 64]) show3Dslice(img_resize) img = img.squeeze(0).data.cpu().numpy() show3Dslice(img) test_index = 0 # img3D[y, x, z] img3D = images[test_index] mask3D = masks[test_index] # show3D(img3D) # show3Dslice(img3D) # 画bbox -- [x1, y1, z1, x2, y2, z2] 绝对 line_thick = 1 bbox3D = result[test_index][0] x1, y1, z1, x2, y2, z2 = bbox3D bbox3D = [y1, x1, z1, y2, x2, z2] bbox3D = np.floor(bbox3D) bbox3D = np.array(bbox3D, dtype=int) # show3Dbbox_img(img3D, bbox3D, 2) bboxInimg = bbox_in_img_for_slice(img3D, bbox3D, line_thick, -2e3) # show3Dslice(bboxInimg) bb_gt = bbox_gt[test_index] # [y1, y2, x1, x2, z1, z2] 绝对 y1, y2, x1, x2, z1, z2 = bb_gt bb_gt = [y1, x1, z1, y2, x2, z2] bb_gt = np.floor(bb_gt) bb_gt = np.array(bb_gt, dtype=int) print(bbox3D) print(bb_gt) bboxInimg = bbox_in_img_for_slice(bboxInimg, bb_gt, line_thick, 1e3) bboxInimg = bboxInimg + mask3D * 2e3 # bbox_in_img = np.hstack((bbox_in_img, img3D)) show3Dslice(bboxInimg) # show3D(bboxInimg) img_for_3D = np.zeros([128, 128, 64], dtype=np.int) img_for_3D = bbox_in_img_for_3D(img_for_3D, bb_gt, line_thick, 100) img_for_3D = bbox_in_img_for_3D(img_for_3D, bbox3D, line_thick, 200) show3D(img_for_3D)
def obj_test(): args = parse_args().parse_args(args=[]) device = torch.device("cpu") input_size = [128, 128, 64] # dataset class_names = ('pnens') # dataset2 = load_from_pkl(r'/data/liyi219/pnens_3D_data/after_dealing/pre_order0_128_128_64_new.pkl') dataset2 = load_from_pkl(r'E:\data\pre_order0_128_128_64_new.pkl') dataset = datanih(dataset2) # load net from models.yolo_v3 import myYOLOv3 anchor_size = anchor_size_3D_try net = myYOLOv3(device, input_size=input_size, num_classes=1, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh, anchor_size=anchor_size) net.load_state_dict( torch.load(os.path.join(args.trained_model + 'yolo_v3_200.pth'), map_location=device)) net.to(device).eval() num_images = len(dataset) result = [] images = [] bbox_gt = [] for index in range(num_images): if index == 4: break print('Testing image {:d}/{:d}....'.format(index + 1, num_images)) img, _, height, width, depth = dataset.pull_item(index) bb = dataset.get_bbox_juedui(index) bbox_gt.append(bb) scale = np.array([[height, width, depth, height, width, depth]]) images.append(img) # to tensor x = img x = x.unsqueeze(0).to(device) t0 = time.time() # forward bboxes, scores, _ = net(x) bboxes = bboxes * scale if not scores.size: print('{} has not bbox'.format(index)) continue best_scores = np.argmax(scores, axis=0) bboxes = bboxes[best_scores] # print(scores[best_scores]) # print([bboxes, best_scores]) result.append([bboxes, best_scores]) print("detection {} time used ".format(index), time.time() - t0, "s") if False: pred_1 = torch.sigmoid(pred_1) pred_1 = pred_1.data.cpu().numpy()[0, 0] img3D = pred_1 vol = mlab.pipeline.volume(mlab.pipeline.scalar_field(img3D), name='3-d ultrasound ') mlab.colorbar(orientation='vertical') mlab.show() show3Dslice(img3D) img_resize = resize3D(img3D, [128, 64, 64]) show3Dslice(img_resize) img = img.squeeze(0).data.cpu().numpy() show3Dslice(img) test_index = 1 img3D = images[test_index] img3D = img3D.squeeze(dim=0) img3D = np.array(img3D) # show3D(img3D) # show3Dslice(img3D) # 画bbox -- [x1, y1, z1, x2, y2, z2] line_thick = 1 bbox3D = result[test_index][0] # print(bbox3D) # print(result[0][1]) bbox3D = np.floor(bbox3D) bbox3D = np.array(bbox3D, dtype=int) # show3Dbbox_img(img3D, bbox3D, 2) bbox_in_img = get_bbox_in_img(img3D, bbox3D, line_thick) bb_gt = bbox_gt[test_index] x1, x2, y1, y2, z1, z2 = bb_gt bb_gt = [x1, y1, z1, x2, y2, z2] bb_gt = np.floor(bb_gt) bb_gt = np.array(bb_gt, dtype=int) print(bbox3D) print(bb_gt) bbox_in_img = get_bbox_in_img(bbox_in_img, bb_gt, line_thick, line_value=2e3) # bbox_in_img = np.hstack((bbox_in_img, img3D)) show3Dslice(bbox_in_img)
# dataloader dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True, collate_fn=detection_collate, num_workers=args.num_workers, pin_memory=True) if args.version == 'yolo_v3': from models.yolo_v3 import myYOLOv3 # anchor_size = MULTI_ANCHOR_SIZE if args.dataset == 'voc' else MULTI_ANCHOR_SIZE_COCO anchor_size = anchor_size_3D_try yolo_net = myYOLOv3(device, input_size=train_size, num_classes=num_classes, trainable=True, anchor_size=anchor_size, hr=hr) print('Let us train yolo_v3 on the %s dataset ......' % (args.dataset)) else: print('Unknown version !!!') exit() model = yolo_net model.to(device).train() # use tfboard if args.tfboard: print('use tensorboard') from torch.utils.tensorboard import SummaryWriter
def train(): args = parse_args() data_dir = args.dataset_root path_to_save = os.path.join(args.save_folder, args.version) os.makedirs(path_to_save, exist_ok=True) hr = False if args.high_resolution: print('use hi-res backbone') hr = True cfg = coco_ab if args.cuda: print('use cuda') cudnn.benchmark = True device = torch.device("cuda") else: device = torch.device("cpu") if args.multi_scale: print('Let us use the multi-scale trick.') input_size = [608, 608] dataset = COCODataset(data_dir=data_dir, img_size=608, transform=SSDAugmentation([608, 608], mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229)), debug=args.debug) else: input_size = cfg['min_dim'] dataset = COCODataset(data_dir=data_dir, img_size=cfg['min_dim'][0], transform=SSDAugmentation(cfg['min_dim'], mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229)), debug=args.debug) # build model if args.version == 'yolo_v2': from models.yolo_v2 import myYOLOv2 total_anchor_size = tools.get_total_anchor_size(name='COCO') yolo_net = myYOLOv2(device, input_size=input_size, num_classes=args.num_classes, trainable=True, anchor_size=total_anchor_size, hr=hr) print('Let us train yolo-v2 on the COCO dataset ......') elif args.version == 'yolo_v3': from models.yolo_v3 import myYOLOv3 total_anchor_size = tools.get_total_anchor_size(multi_level=True, name='COCO') yolo_net = myYOLOv3(device, input_size=input_size, num_classes=args.num_classes, trainable=True, anchor_size=total_anchor_size, hr=hr) print('Let us train yolo-v3 on the COCO dataset ......') elif args.version == 'tiny_yolo_v2': from models.tiny_yolo_v2 import YOLOv2tiny total_anchor_size = tools.get_total_anchor_size(name='COCO') yolo_net = YOLOv2tiny(device, input_size=input_size, num_classes=args.num_classes, trainable=True, anchor_size=total_anchor_size, hr=hr) print('Let us train tiny-yolo-v2 on the COCO dataset ......') elif args.version == 'tiny_yolo_v3': from models.tiny_yolo_v3 import YOLOv3tiny total_anchor_size = tools.get_total_anchor_size(multi_level=True, name='COCO') yolo_net = YOLOv3tiny(device, input_size=input_size, num_classes=args.num_classes, trainable=True, anchor_size=total_anchor_size, hr=hr) print('Let us train tiny-yolo-v3 on the COCO dataset ......') else: print('Unknown version !!!') exit() print("Setting Arguments.. : ", args) print("----------------------------------------------------------") print('Loading the MSCOCO dataset...') print('Training model on:', dataset.name) print('The dataset size:', len(dataset)) print("----------------------------------------------------------") # use tfboard if args.tfboard: print('use tensorboard') from torch.utils.tensorboard import SummaryWriter c_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time.time())) log_path = os.path.join('log/coco/', args.version, c_time) os.makedirs(log_path, exist_ok=True) writer = SummaryWriter(log_path) print('Let us train yolo-v2 on the MSCOCO dataset ......') model = yolo_net model.to(device).train() dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True, collate_fn=detection_collate, num_workers=args.num_workers) evaluator = COCOAPIEvaluator(data_dir=data_dir, img_size=cfg['min_dim'], device=device, transform=BaseTransform(cfg['min_dim'], mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229))) # optimizer setup base_lr = args.lr tmp_lr = base_lr optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) max_epoch = cfg['max_epoch'] epoch_size = len(dataset) // args.batch_size # start training loop t0 = time.time() for epoch in range(max_epoch): # use cos lr if args.cos and epoch > 20 and epoch <= max_epoch - 20: # use cos lr tmp_lr = 0.00001 + 0.5 * (base_lr - 0.00001) * ( 1 + math.cos(math.pi * (epoch - 20) * 1. / (max_epoch - 20))) set_lr(optimizer, tmp_lr) elif args.cos and epoch > max_epoch - 20: tmp_lr = 0.00001 set_lr(optimizer, tmp_lr) # use step lr else: if epoch in cfg['lr_epoch']: tmp_lr = tmp_lr * 0.1 set_lr(optimizer, tmp_lr) # COCO evaluation if (epoch + 1) % args.eval_epoch == 0: model.trainable = False model.set_grid(cfg['min_dim']) # evaluate ap50_95, ap50 = evaluator.evaluate(model) print('ap50 : ', ap50) print('ap50_95 : ', ap50_95) # convert to training mode. model.trainable = True model.set_grid(input_size) model.train() if args.tfboard: writer.add_scalar('val/COCOAP50', ap50, epoch + 1) writer.add_scalar('val/COCOAP50_95', ap50_95, epoch + 1) for iter_i, (images, targets) in enumerate(dataloader): # WarmUp strategy for learning rate if not args.no_warm_up: if epoch < args.wp_epoch: # tmp_lr = base_lr * pow((iter_i+epoch*epoch_size)*1. / (args.wp_epoch*epoch_size), 4) tmp_lr = 1e-6 + (base_lr - 1e-6) * ( iter_i + epoch * epoch_size) / (epoch_size * (args.wp_epoch)) set_lr(optimizer, tmp_lr) elif epoch == args.wp_epoch and iter_i == 0: tmp_lr = base_lr set_lr(optimizer, tmp_lr) targets = [label.tolist() for label in targets] if args.version == 'yolo_v2' or args.version == 'tiny_yolo_v2': targets = tools.gt_creator(input_size, yolo_net.stride, targets, name='COCO') elif args.version == 'yolo_v3' or args.version == 'tiny_yolo_v3': targets = tools.multi_gt_creator(input_size, yolo_net.stride, targets, name='COCO') # to device images = images.to(device) targets = torch.tensor(targets).float().to(device) # forward and loss conf_loss, cls_loss, txtytwth_loss, total_loss = model( images, target=targets) # backprop total_loss.backward() optimizer.step() optimizer.zero_grad() if args.tfboard: # viz loss writer.add_scalar('object loss', conf_loss.item(), iter_i + epoch * epoch_size) writer.add_scalar('class loss', cls_loss.item(), iter_i + epoch * epoch_size) writer.add_scalar('local loss', txtytwth_loss.item(), iter_i + epoch * epoch_size) if iter_i % 10 == 0: t1 = time.time() print( '[Epoch %d/%d][Iter %d/%d][lr %.6f]' '[Loss: obj %.2f || cls %.2f || bbox %.2f || total %.2f || size %d || time: %.2f]' % (epoch + 1, max_epoch, iter_i, epoch_size, tmp_lr, conf_loss.item(), cls_loss.item(), txtytwth_loss.item(), total_loss.item(), input_size[0], t1 - t0), flush=True) t0 = time.time() # multi-scale trick if iter_i % 10 == 0 and iter_i > 0 and args.multi_scale: # ms_ind = random.sample(ms_inds, 1)[0] # input_size = cfg['multi_scale'][int(ms_ind)] size = random.randint(10, 19) * 32 input_size = [size, size] model.set_grid(input_size) # change input dim # But this operation will report bugs when we use more workers in data loader, so I have to use 0 workers. # I don't know how to make it suit more workers, and I'm trying to solve this question. dataloader.dataset.reset_transform( SSDAugmentation(input_size, mean=(0.406, 0.456, 0.485), std=(0.225, 0.224, 0.229))) if (epoch + 1) % 10 == 0: print('Saving state, epoch:', epoch + 1) torch.save( model.state_dict(), os.path.join(path_to_save, args.version + '_' + repr(epoch + 1) + '.pth'))
def see_feature_segment(): args = parse_args().parse_args(args=[]) device = torch.device("cpu") input_size = [128, 128, 64] # dataset num_classes = 1 class_names = ('pnens') # dataset2 = load_from_pkl(r'/data/liyi219/pnens_3D_data/after_dealing/pre_order0_128_128_64_new.pkl') dataset2 = load_from_pkl(r'E:\data\pre_order0_128_128_64_new.pkl') dataset = datanih(dataset2) # load net from models.yolo_v3 import myYOLOv3 anchor_size = anchor_size_3D_try net = myYOLOv3(device, input_size=input_size, num_classes=1, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh, anchor_size=anchor_size) net.load_state_dict( torch.load(os.path.join(args.trained_model + 'yolo_v3_200.pth'), map_location=device)) net.to(device).eval() num_images = len(dataset) result = [] images = [] for index in range(num_images): if index == 1: break print('Testing image {:d}/{:d}....'.format(index + 1, num_images)) img, _, height, width, depth = dataset.pull_item(index) scale = np.array([[height, width, depth, height, width, depth]]) images.append(img) # to tensor x = img x = x.unsqueeze(0).to(device) t0 = time.time() # forward bboxes, scores, _ = net(x) pred_1, pred_2, pred_3 = net.get_feature_map(x) bboxes = bboxes * scale if not scores.size: print('{} has not bbox'.format(index)) continue best_scores = np.argmax(scores, axis=0) bboxes = bboxes[best_scores] result.append([bboxes, scores]) print("detection {} time used ".format(index), time.time() - t0, "s") img3D = images[0] img3D = img3D.squeeze(dim=0) img3D = np.array(img3D) pred_1 = torch.sigmoid(pred_1) pred_1 = pred_1.data.cpu().numpy()[0, 0] show3D(pred_1) img_resize = resize3D(pred_1, [128, 128, 64]) # show3Dslice(img_resize) img_and_fearturemap = np.hstack((img_resize, img3D)) show3Dslice(img_and_fearturemap)