def test_net(): net=Resnet18_8s(ver_dim=vote_num*2, seg_dim=2) net=NetWrapper(net) net=DataParallel(net).cuda() optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr']) model_dir=os.path.join(cfg.MODEL_DIR,train_cfg['model_name']) motion_model=train_cfg['motion_model'] print('motion state {}'.format(motion_model)) # if args.test_model: begin_epoch=load_model(net.module.net, optimizer, model_dir, args.load_epoch) print('testing normal linemod ...') image_db = LineModImageDB(args.linemod_cls,has_render_set=False, has_fuse_set=False) test_db = image_db.test_real_set+image_db.val_real_set test_set = LineModDatasetRealAug(test_db, cfg.LINEMOD, vote_type, augment=False, use_motion=motion_model) test_sampler = SequentialSampler(test_set) test_batch_sampler = ImageSizeBatchSampler(test_sampler, train_cfg['test_batch_size'], False) test_loader = DataLoader(test_set, batch_sampler=test_batch_sampler, num_workers=0) args.use_test_set = True prefix='test' val(net, test_loader, begin_epoch, prefix, use_motion=motion_model) if args.linemod_cls in cfg.occ_linemod_cls_names: print('testing occluded linemod ...') occ_image_db = OcclusionLineModImageDB(args.linemod_cls) occ_test_db = occ_image_db.test_real_set occ_test_set = LineModDatasetRealAug(occ_test_db, cfg.OCCLUSION_LINEMOD, vote_type, augment=False, use_motion=motion_model) occ_test_sampler = SequentialSampler(occ_test_set) occ_test_batch_sampler = ImageSizeBatchSampler(occ_test_sampler, train_cfg['test_batch_size'], False) occ_test_loader = DataLoader(occ_test_set, batch_sampler=occ_test_batch_sampler, num_workers=0) prefix='occ_test' if args.use_test_set else 'occ_val' val(net, occ_test_loader, begin_epoch, prefix, use_motion=motion_model)
def train_net(): #seg_dim 为啥两张,kp 个数,×2,? 一张一个x坐标,一张一个y坐标,不是一张图片估计 一个关键点吗 net = Resnet18_8s(ver_dim=vote_num * 2, seg_dim=2) #前传 误差定义 评价 net = NetWrapper(net) #多卡 #net=DataParallel(net,device_ids=[0,1]).cuda() net = DataParallel(net, device_ids=[0]).cuda() #优化方法 optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr']) model_dir = os.path.join(cfg.MODEL_DIR, train_cfg['model_name']) motion_model = train_cfg['motion_model'] print('motion state {}'.format(motion_model)) if args.test_model: begin_epoch = load_model(net.module.net, optimizer, model_dir, args.load_epoch) if args.normal: print('testing normal linemod ...') image_db = LineModImageDB(args.linemod_cls, has_render_set=False, has_fuse_set=False) test_db = image_db.test_real_set + image_db.val_real_set test_set = LineModDatasetRealAug(test_db, cfg.LINEMOD, vote_type, augment=False, use_motion=motion_model) test_sampler = SequentialSampler(test_set) test_batch_sampler = ImageSizeBatchSampler( test_sampler, train_cfg['test_batch_size'], False) test_loader = DataLoader(test_set, batch_sampler=test_batch_sampler, num_workers=0) prefix = 'test' if args.use_test_set else 'val' val(net, test_loader, begin_epoch, prefix, use_motion=motion_model) if args.occluded and args.linemod_cls in cfg.occ_linemod_cls_names: print('testing occluded linemod ...') occ_image_db = OcclusionLineModImageDB(args.linemod_cls) occ_test_db = occ_image_db.test_real_set occ_test_set = LineModDatasetRealAug(occ_test_db, cfg.OCCLUSION_LINEMOD, vote_type, augment=False, use_motion=motion_model) occ_test_sampler = SequentialSampler(occ_test_set) occ_test_batch_sampler = ImageSizeBatchSampler( occ_test_sampler, train_cfg['test_batch_size'], False) occ_test_loader = DataLoader(occ_test_set, batch_sampler=occ_test_batch_sampler, num_workers=0) prefix = 'occ_test' if args.use_test_set else 'occ_val' val(net, occ_test_loader, begin_epoch, prefix, use_motion=motion_model) if args.truncated: print('testing truncated linemod ...') trun_image_db = TruncatedLineModImageDB(args.linemod_cls) print(len(trun_image_db.set)) trun_image_set = LineModDatasetRealAug(trun_image_db.set, cfg.LINEMOD, vote_type, augment=False, use_intrinsic=True, use_motion=motion_model) trun_test_sampler = SequentialSampler(trun_image_set) trun_test_batch_sampler = ImageSizeBatchSampler( trun_test_sampler, train_cfg['test_batch_size'], False) trun_test_loader = DataLoader( trun_image_set, batch_sampler=trun_test_batch_sampler, num_workers=0) prefix = 'trun_test' val(net, trun_test_loader, begin_epoch, prefix, True, use_motion=motion_model) else: # //train from finetune,or train from 0 # "resume":true, # //this para no use # "finetune":false, begin_epoch = 0 # train from finetune if train_cfg['resume']: begin_epoch = load_model(net.module.net, optimizer, model_dir) #cat already render # image_db = LineModImageDB(args.linemod_cls, # has_fuse_set=train_cfg['use_fuse'], # has_render_set=True) image_db = LineModImageDB(args.linemod_cls, has_fuse_set=train_cfg['use_fuse'], has_render_set=False) train_db = [] #三类训练数据吗,render,real,fuse?? 阶段数据呢,,遮挡数据呢 train_db += image_db.render_set if train_cfg['use_real_train']: train_db += image_db.train_real_set if train_cfg['use_fuse']: train_db += image_db.fuse_set train_set = LineModDatasetRealAug(train_db, cfg.LINEMOD, vote_type, augment=True, cfg=train_cfg['aug_cfg'], use_motion=motion_model) train_sampler = RandomSampler(train_set) train_batch_sampler = ImageSizeBatchSampler( train_sampler, train_cfg['train_batch_size'], False, cfg=train_cfg['aug_cfg']) train_loader = DataLoader(train_set, batch_sampler=train_batch_sampler, num_workers=12) val_db = image_db.val_real_set val_set = LineModDatasetRealAug(val_db, cfg.LINEMOD, vote_type, augment=False, cfg=train_cfg['aug_cfg'], use_motion=motion_model) val_sampler = SequentialSampler(val_set) val_batch_sampler = ImageSizeBatchSampler(val_sampler, train_cfg['test_batch_size'], False, cfg=train_cfg['aug_cfg']) val_loader = DataLoader(val_set, batch_sampler=val_batch_sampler, num_workers=12) if args.linemod_cls in cfg.occ_linemod_cls_names: occ_image_db = OcclusionLineModImageDB(args.linemod_cls) occ_val_db = occ_image_db.test_real_set[:len(occ_image_db. test_real_set) // 2] occ_val_set = LineModDatasetRealAug(occ_val_db, cfg.OCCLUSION_LINEMOD, vote_type, augment=False, cfg=train_cfg['aug_cfg'], use_motion=motion_model) occ_val_sampler = SequentialSampler(occ_val_set) occ_val_batch_sampler = ImageSizeBatchSampler( occ_val_sampler, train_cfg['test_batch_size'], False, cfg=train_cfg['aug_cfg']) occ_val_loader = DataLoader(occ_val_set, batch_sampler=occ_val_batch_sampler, num_workers=12) for epoch in range(begin_epoch, train_cfg['epoch_num']): adjust_learning_rate(optimizer, epoch, train_cfg['lr_decay_rate'], train_cfg['lr_decay_epoch']) print("xx") train(net, optimizer, train_loader, epoch) val(net, val_loader, epoch, use_motion=motion_model) if args.linemod_cls in cfg.occ_linemod_cls_names: val(net, occ_val_loader, epoch, 'occ_val', use_motion=motion_model) save_model(net.module.net, optimizer, epoch, model_dir)
sys.path.append('.') from lib.utils.data_utils import LineModImageDB,LineModModelDB,Projector from lib.datasets.linemod_dataset import LineModDatasetRealAug,ImageSizeBatchSampler,VotingType from lib.ransac_voting_gpu_layer.ransac_voting_gpu import ransac_voting_layer from torch.utils.data import RandomSampler,DataLoader from lib.utils.draw_utils import pts_to_img_pts from lib.utils.evaluation_utils import pnp import random image_db = LineModImageDB('duck', has_ro_set=False, has_ra_set=False, has_plane_set=False, has_render_set=False, has_ms_set=False,has_fuse_set=False) random.shuffle(image_db.real_set) dataset = LineModDatasetRealAug(image_db.real_set[:5], data_prefix=image_db.linemod_dir, vote_type=VotingType.Extreme, augment=False) sampler = RandomSampler(dataset) batch_sampler = ImageSizeBatchSampler(sampler, 5, False) loader = DataLoader(dataset, batch_sampler=batch_sampler, num_workers=8) modeldb=LineModModelDB() camera_matrix=Projector().intrinsic_matrix['linemod'].astype(np.float32) for i, data in enumerate(loader): rgb, mask, vertex, vertex_weight, pose, gt_corners = data pts2d=gt_corners[0].numpy()[:,:2].astype(np.float32) pts3d=modeldb.get_extreme_3d('duck') pts3d=np.concatenate([pts3d,modeldb.get_centers_3d('duck')[None,:]],0).astype(np.float32) wgt2d=np.zeros([pts2d.shape[0],3]).astype(np.float32) wgt2d[:,(0,2)]=1.0 for k in range(pts2d.shape[0]): if np.random.random()<0.5: scale = np.random.uniform(1, 8)
def train_net(): net = Resnet18_8s_modified(ver_dim=vote_num * 2, seg_dim=2) #net=NetWrapper(net) net = MappingNetWrapper(net) net = DataParallel(net).cuda() optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr']) model_dir = os.path.join(cfg.MODEL_DIR, train_cfg['model_name']) motion_model = train_cfg['motion_model'] print('motion state {}'.format(motion_model)) if args.test_model: begin_epoch = load_model(net.module.net, optimizer, model_dir, args.load_epoch) if args.normal: print('testing normal linemod ...') image_db = LineModImageDB(args.linemod_cls, has_render_set=False, has_fuse_set=False) test_db = image_db.test_real_set + image_db.val_real_set test_set = LineModDatasetAug(test_db, cfg.LINEMOD, vote_type, augment=False, use_motion=motion_model) test_sampler = SequentialSampler(test_set) test_batch_sampler = ImageSizeBatchSampler( test_sampler, train_cfg['test_batch_size'], False) test_loader = DataLoader(test_set, batch_sampler=test_batch_sampler, num_workers=0) prefix = 'test' if args.use_test_set else 'val' val(net, test_loader, begin_epoch, prefix, use_motion=motion_model) # if args.occluded and args.linemod_cls in cfg.occ_linemod_cls_names: # print('testing occluded linemod ...') # occ_image_db = OcclusionLineModImageDB(args.linemod_cls) # occ_test_db = occ_image_db.test_real_set # occ_test_set = LineModDatasetAug(occ_test_db, # cfg.OCCLUSION_LINEMOD, # vote_type, # augment=False, # use_motion=motion_model) # occ_test_sampler = SequentialSampler(occ_test_set) # occ_test_batch_sampler = ImageSizeBatchSampler( # occ_test_sampler, train_cfg['test_batch_size'], False) # occ_test_loader = DataLoader(occ_test_set, # batch_sampler=occ_test_batch_sampler, # num_workers=0) # prefix = 'occ_test' if args.use_test_set else 'occ_val' # val(net, # occ_test_loader, # begin_epoch, # prefix, # use_motion=motion_model) # if args.truncated: # print('testing truncated linemod ...') # trun_image_db = TruncatedLineModImageDB(args.linemod_cls) # print(len(trun_image_db.set)) # trun_image_set = LineModDatasetRealAug(trun_image_db.set, # cfg.LINEMOD, # vote_type, # augment=False, # use_intrinsic=True, # use_motion=motion_model) # trun_test_sampler = SequentialSampler(trun_image_set) # trun_test_batch_sampler = ImageSizeBatchSampler( # trun_test_sampler, train_cfg['test_batch_size'], False) # trun_test_loader = DataLoader( # trun_image_set, # batch_sampler=trun_test_batch_sampler, # num_workers=0) # prefix = 'trun_test' # val(net, # trun_test_loader, # begin_epoch, # prefix, # True, # use_motion=motion_model) else: begin_epoch = 0 if train_cfg['resume']: begin_epoch = load_model(net.module.net, optimizer, model_dir) image_db = LineModImageDB( args.linemod_cls, has_fuse_set=False, #train_cfg['use_fuse'], has_render_set=False) #True) # train_db=[] # train_db+=image_db.render_set # if train_cfg['use_real_train']: # train_db+=image_db.train_real_set # if train_cfg['use_fuse']: # train_db+=image_db.fuse_set train_db = image_db.train_real_set train_set = LineModDatasetAug(train_db, cfg.LINEMOD, vote_type, augment=True, cfg=train_cfg['aug_cfg'], use_motion=motion_model) train_sampler = RandomSampler(train_set) train_batch_sampler = ImageSizeBatchSampler( train_sampler, train_cfg['train_batch_size'], False, cfg=train_cfg['aug_cfg']) train_loader = DataLoader(train_set, batch_sampler=train_batch_sampler, num_workers=12) val_db = image_db.val_real_set val_set = LineModDatasetAug(val_db, cfg.LINEMOD, vote_type, augment=False, cfg=train_cfg['aug_cfg'], use_motion=motion_model) val_sampler = SequentialSampler(val_set) val_batch_sampler = ImageSizeBatchSampler(val_sampler, train_cfg['test_batch_size'], False, cfg=train_cfg['aug_cfg']) val_loader = DataLoader(val_set, batch_sampler=val_batch_sampler, num_workers=12) # if args.linemod_cls in cfg.occ_linemod_cls_names: # occ_image_db=OcclusionLineModImageDB(args.linemod_cls) # occ_val_db=occ_image_db.test_real_set[:len(occ_image_db.test_real_set)//2] # occ_val_set = LineModDatasetRealAug(occ_val_db, cfg.OCCLUSION_LINEMOD, vote_type, augment=False, cfg=train_cfg['aug_cfg'], use_motion=motion_model) # occ_val_sampler = SequentialSampler(occ_val_set) # occ_val_batch_sampler = ImageSizeBatchSampler(occ_val_sampler, train_cfg['test_batch_size'], False, cfg=train_cfg['aug_cfg']) # occ_val_loader = DataLoader(occ_val_set, batch_sampler=occ_val_batch_sampler, num_workers=12) for epoch in range(begin_epoch, train_cfg['epoch_num']): adjust_learning_rate(optimizer, epoch, train_cfg['lr_decay_rate'], train_cfg['lr_decay_epoch']) train(net, optimizer, train_loader, epoch) val(net, val_loader, epoch, use_motion=motion_model) # if args.linemod_cls in cfg.occ_linemod_cls_names: # val(net, occ_val_loader, epoch, 'occ_val',use_motion=motion_model) save_model(net.module.net, optimizer, epoch, model_dir)
image_db = LineModImageDB(args.linemod_cls, has_fuse_set=train_cfg['use_fuse'], has_render_set=True) train_db=[] train_db+=image_db.render_set if train_cfg['use_real_train']: train_db+=image_db.train_real_set if train_cfg['use_fuse']: train_db+=image_db.fuse_set train_set = LineModDatasetRealAug(train_db, cfg.LINEMOD, vote_type, augment=True, cfg=train_cfg['aug_cfg'], use_motion=motion_model) train_sampler = RandomSampler(train_set) train_batch_sampler = ImageSizeBatchSampler(train_sampler, train_cfg['train_batch_size'], False, cfg=train_cfg['aug_cfg']) train_loader = DataLoader(train_set, batch_sampler=train_batch_sampler, num_workers=12) val_db=image_db.val_real_set val_set = LineModDatasetRealAug(val_db, cfg.LINEMOD, vote_type, augment=False, cfg=train_cfg['aug_cfg'], use_motion=motion_model) val_sampler = SequentialSampler(val_set) val_batch_sampler = ImageSizeBatchSampler(val_sampler, train_cfg['test_batch_size'], False, cfg=train_cfg['aug_cfg']) val_loader = DataLoader(val_set, batch_sampler=val_batch_sampler, num_workers=12) if args.linemod_cls in cfg.occ_linemod_cls_names: occ_image_db=OcclusionLineModImageDB(args.linemod_cls) occ_val_db=occ_image_db.test_real_set[:len(occ_image_db.test_real_set)//2] occ_val_set = LineModDatasetRealAug(occ_val_db, cfg.OCCLUSION_LINEMOD, vote_type, augment=False, cfg=train_cfg['aug_cfg'], use_motion=motion_model) occ_val_sampler = SequentialSampler(occ_val_set) occ_val_batch_sampler = ImageSizeBatchSampler(occ_val_sampler, train_cfg['test_batch_size'], False, cfg=train_cfg['aug_cfg']) occ_val_loader = DataLoader(occ_val_set, batch_sampler=occ_val_batch_sampler, num_workers=12)
def train_net(): net = Resnet18_8s(ver_dim=vote_num * 2, seg_dim=2) net = NetWrapper(net) net = DataParallel(net).cuda() optimizer = optim.Adam(net.parameters(), lr=train_cfg["lr"]) model_dir = os.path.join(cfg.MODEL_DIR, train_cfg["model_name"]) motion_model = train_cfg["motion_model"] print("motion state {}".format(motion_model)) if args.test_model: begin_epoch = load_model(net.module.net, optimizer, model_dir, args.load_epoch) if args.normal: print("testing normal linemod ...") image_db = LineModImageDB(args.linemod_cls, has_render_set=False, has_fuse_set=False) test_db = image_db.test_real_set + image_db.val_real_set test_set = LineModDatasetRealAug(test_db, cfg.LINEMOD, vote_type, augment=False, use_motion=motion_model) test_sampler = SequentialSampler(test_set) test_batch_sampler = ImageSizeBatchSampler( test_sampler, train_cfg["test_batch_size"], False) test_loader = DataLoader(test_set, batch_sampler=test_batch_sampler, num_workers=0) prefix = "test" if args.use_test_set else "val" val(net, test_loader, begin_epoch, prefix, use_motion=motion_model) if args.occluded and args.linemod_cls in cfg.occ_linemod_cls_names: print("testing occluded linemod ...") occ_image_db = OcclusionLineModImageDB(args.linemod_cls) occ_test_db = occ_image_db.test_real_set occ_test_set = LineModDatasetRealAug( occ_test_db, cfg.OCCLUSION_LINEMOD, vote_type, augment=False, use_motion=motion_model, ) occ_test_sampler = SequentialSampler(occ_test_set) occ_test_batch_sampler = ImageSizeBatchSampler( occ_test_sampler, train_cfg["test_batch_size"], False) occ_test_loader = DataLoader(occ_test_set, batch_sampler=occ_test_batch_sampler, num_workers=0) prefix = "occ_test" if args.use_test_set else "occ_val" val(net, occ_test_loader, begin_epoch, prefix, use_motion=motion_model) if args.truncated: print("testing truncated linemod ...") trun_image_db = TruncatedLineModImageDB(args.linemod_cls) print(len(trun_image_db.set)) trun_image_set = LineModDatasetRealAug( trun_image_db.set, cfg.LINEMOD, vote_type, augment=False, use_intrinsic=True, use_motion=motion_model, ) trun_test_sampler = SequentialSampler(trun_image_set) trun_test_batch_sampler = ImageSizeBatchSampler( trun_test_sampler, train_cfg["test_batch_size"], False) trun_test_loader = DataLoader( trun_image_set, batch_sampler=trun_test_batch_sampler, num_workers=0) prefix = "trun_test" val( net, trun_test_loader, begin_epoch, prefix, True, use_motion=motion_model, ) else: begin_epoch = 0 if train_cfg["resume"]: begin_epoch = load_model(net.module.net, optimizer, model_dir) image_db = LineModImageDB(args.linemod_cls, has_fuse_set=train_cfg["use_fuse"], has_render_set=False) train_db = [] train_db += image_db.render_set if train_cfg["use_real_train"]: train_db += image_db.train_real_set if train_cfg["use_fuse"]: train_db += image_db.fuse_set train_set = LineModDatasetRealAug( train_db, cfg.LINEMOD, vote_type, augment=True, cfg=train_cfg["aug_cfg"], use_motion=motion_model, ) train_sampler = RandomSampler(train_set) train_batch_sampler = ImageSizeBatchSampler( train_sampler, train_cfg["train_batch_size"], False, cfg=train_cfg["aug_cfg"], ) train_loader = DataLoader(train_set, batch_sampler=train_batch_sampler, num_workers=12) val_db = image_db.val_real_set val_set = LineModDatasetRealAug( val_db, cfg.LINEMOD, vote_type, augment=False, cfg=train_cfg["aug_cfg"], use_motion=motion_model, ) val_sampler = SequentialSampler(val_set) val_batch_sampler = ImageSizeBatchSampler(val_sampler, train_cfg["test_batch_size"], False, cfg=train_cfg["aug_cfg"]) val_loader = DataLoader(val_set, batch_sampler=val_batch_sampler, num_workers=12) if args.use_occlussion and args.linemod_cls in cfg.occ_linemod_cls_names: occ_image_db = OcclusionLineModImageDB(args.linemod_cls) occ_val_db = occ_image_db.test_real_set[:len(occ_image_db. test_real_set) // 2] occ_val_set = LineModDatasetRealAug( occ_val_db, cfg.OCCLUSION_LINEMOD, vote_type, augment=False, cfg=train_cfg["aug_cfg"], use_motion=motion_model, ) occ_val_sampler = SequentialSampler(occ_val_set) occ_val_batch_sampler = ImageSizeBatchSampler( occ_val_sampler, train_cfg["test_batch_size"], False, cfg=train_cfg["aug_cfg"], ) occ_val_loader = DataLoader(occ_val_set, batch_sampler=occ_val_batch_sampler, num_workers=12) for epoch in range(begin_epoch, train_cfg["epoch_num"]): adjust_learning_rate( optimizer, epoch, train_cfg["lr_decay_rate"], train_cfg["lr_decay_epoch"], ) train(net, optimizer, train_loader, epoch) val(net, val_loader, epoch, use_motion=motion_model) if args.use_occlussion and args.linemod_cls in cfg.occ_linemod_cls_names: val(net, occ_val_loader, epoch, "occ_val", use_motion=motion_model) save_model(net.module.net, optimizer, epoch, model_dir)
def train_net(): # mp.spawn(demo_basic, # args=(4,), # nprocs=4, # join=True) tf_dir = './runs/' + train_cfg['exp_name'] writer = SummaryWriter(log_dir=tf_dir) Path("/home/gerard/myPvnet/pvnet/{}".format(train_cfg["exp_name"])).mkdir( parents=True, exist_ok=True) model_dir = os.path.join(cfg.MODEL_DIR, train_cfg['model_name']) imNet = ImageUNet(ver_dim=(vote_num * 2), seg_dim=2) estNet = EstimateUNet(ver_dim=(vote_num * 2), seg_dim=2) net = NetWrapper(imNet, estNet) net = DataParallel(net).cuda() # if train_cfg['exp_name'] == 'AE_ape': # model_Dir=os.path.join(cfg.MODEL_DIR,'ape_linemod_train_GE') # imNet=load_pretrained_imNet(ImageUNet(ver_dim=(vote_num*2), seg_dim=2), model_Dir, epoch=0) # elif train_cfg['exp_name'] == 'GE_ape': # model_Dir=os.path.join(cfg.MODEL_DIR,'ape_linemod_train_AE') # # estNet=load_pretrained_estNet(EstimateUNet(ver_dim=(vote_num*2), seg_dim=2), model_Dir, epoch=25) # elif train_cfg['exp_name'] == 'AE_GE_ape': # model_Dir=os.path.join(cfg.PVModelDir)['net']) # PVNet = PVNet.cuda(rank) # PVNet=DistributedDataParallel(PVNet, device_ids=[rank]) PVModelDir = '/home/gerard/baseline_models/{}_baseline/199.pth'.format( train_cfg['object']) PVNet = PVnet(ver_dim=vote_num * 2, seg_dim=2) PVNet.load_state_dict(torch.load(PVModelDir)['net']) PVNet = PVNet.half() PVNet = DataParallel(PVNet).cuda() randomCropping = RandomScaleCrop() optimizer = optim.Adam(net.parameters(), lr=train_cfg['lr']) motion_model = train_cfg['motion_model'] print('motion state {}'.format(motion_model)) for param_group in optimizer.param_groups: lr = param_group['lr'] image_db = LineModImageDB(args.linemod_cls, has_render_set=False, has_fuse_set=False) test_db = image_db.test_real_set + image_db.val_real_set test_set = LineModDatasetRealAug(test_db, randomCropping, cfg.LINEMOD, vote_type, augment=False, use_motion=motion_model) test_sampler = SequentialSampler(test_set) test_batch_sampler = ImageSizeBatchSampler(test_sampler, train_cfg['test_batch_size'], False) test_loader = DataLoader(test_set, batch_sampler=test_batch_sampler, num_workers=0) prefix = 'test' #if args.use_test_set else 'val' if args.test_model: begin_epoch = load_model(net.module.imNet, net.module.estNet, optimizer, model_dir, args.load_epoch) if args.normal: print('testing normal linemod ...') # image_db = LineModImageDB(args.linemod_cls,has_render_set=False, # has_fuse_set=False) # test_db = image_db.test_real_set+image_db.val_real_set # test_set = LineModDatasetRealAug(test_db, randomCropping, cfg.LINEMOD, vote_type, augment=False, use_motion=motion_model) # test_sampler = SequentialSampler(test_set) # test_batch_sampler = ImageSizeBatchSampler(test_sampler, train_cfg['test_batch_size'], False) # test_loader = DataLoader(test_set, batch_sampler=test_batch_sampler, num_workers=0) val_db = image_db.test_real_set + image_db.val_real_set val_set = LineModDatasetRealAug(val_db, randomCropping, cfg.LINEMOD, vote_type, augment=False, cfg=train_cfg['aug_cfg'], use_motion=motion_model) val_sampler = SequentialSampler(val_set) val_batch_sampler = ImageSizeBatchSampler( val_sampler, train_cfg['test_batch_size'], False, cfg=train_cfg['aug_cfg']) val_loader = DataLoader(val_set, batch_sampler=val_batch_sampler, num_workers=16) print('test with val loader') _, _, _, _, _, _, _, _, _ = val(net, PVNet, val_loader, begin_epoch, lr, writer, use_motion=motion_model) prefix = 'test' if args.use_test_set else 'val' print('test with test_loader') _, _, _, _, _, _, _, _, _ = val(net, PVNet, test_loader, begin_epoch, lr, writer, use_motion=motion_model) if args.occluded and args.linemod_cls in cfg.occ_linemod_cls_names: print('testing occluded linemod ...') occ_image_db = OcclusionLineModImageDB(args.linemod_cls) occ_test_db = occ_image_db.test_real_set occ_test_set = LineModDatasetRealAug(occ_test_db, randomCropping, cfg.OCCLUSION_LINEMOD, vote_type, augment=False, use_motion=motion_model) occ_test_sampler = SequentialSampler(occ_test_set) occ_test_batch_sampler = ImageSizeBatchSampler( occ_test_sampler, train_cfg['test_batch_size'], False) occ_test_loader = DataLoader(occ_test_set, batch_sampler=occ_test_batch_sampler, num_workers=0) prefix = 'occ_test' if args.use_test_set else 'occ_val' _, _, _, _, _, _, _, _, _ = val(net, PVNet, occ_test_loader, begin_epoch, lr, writer, prefix, use_motion=motion_model) if args.truncated: print('testing truncated linemod ...') trun_image_db = TruncatedLineModImageDB(args.linemod_cls) print(len(trun_image_db.set)) trun_image_set = LineModDatasetRealAug(trun_image_db.set, cfg.LINEMOD, vote_type, randomCropping, augment=False, use_intrinsic=True, use_motion=motion_model) trun_test_sampler = SequentialSampler(trun_image_set) trun_test_batch_sampler = ImageSizeBatchSampler( trun_test_sampler, train_cfg['test_batch_size'], False) trun_test_loader = DataLoader( trun_image_set, batch_sampler=trun_test_batch_sampler, num_workers=16) prefix = 'trun_test' _, _, _, _, _, _, _, _, _ = val(net, PVNet, trun_test_loader, begin_epoch, lr, writer, prefix, True, use_motion=motion_model) else: begin_epoch = 0 if train_cfg['resume']: begin_epoch = load_model(net.module.imNet, net.module.estNet, optimizer, model_dir) # reset learning rate for param_group in optimizer.param_groups: param_group['lr'] = train_cfg['lr'] lr = param_group['lr'] image_db = LineModImageDB(args.linemod_cls, has_fuse_set=train_cfg['use_fuse'], has_render_set=True) train_db = [] train_db += image_db.render_set if train_cfg['use_real_train']: train_db += image_db.train_real_set if train_cfg['use_fuse']: train_db += image_db.fuse_set train_set = LineModDatasetRealAug(train_db, randomCropping, cfg.LINEMOD, vote_type, augment=True, cfg=train_cfg['aug_cfg'], use_motion=motion_model) train_sampler = RandomSampler(train_set) # train_sampler = torch.utils.data.distributed.DistributedSampler( # train_set, # num_replicas=world_size, # rank=rank # ) # train_batch_sampler = ImageSizeBatchSampler(train_sampler, int(train_cfg['train_batch_size']/world_size), False, cfg=train_cfg['aug_cfg']) train_batch_sampler = ImageSizeBatchSampler( train_sampler, train_cfg['train_batch_size'], False, cfg=train_cfg['aug_cfg']) # train_loader = DataLoader(train_set, batch_sampler=train_batch_sampler, shuffle=False, num_workers=16, pin_memory=True) train_loader = DataLoader(train_set, batch_sampler=train_batch_sampler, num_workers=16) val_db = image_db.test_real_set + image_db.val_real_set val_set = LineModDatasetRealAug(val_db, randomCropping, cfg.LINEMOD, vote_type, augment=False, cfg=train_cfg['aug_cfg'], use_motion=motion_model) val_sampler = SequentialSampler(val_set) val_batch_sampler = ImageSizeBatchSampler(val_sampler, train_cfg['test_batch_size'], False, cfg=train_cfg['aug_cfg']) val_loader = DataLoader(val_set, batch_sampler=val_batch_sampler, num_workers=16) if args.linemod_cls in cfg.occ_linemod_cls_names: occ_image_db = OcclusionLineModImageDB(args.linemod_cls) occ_val_db = occ_image_db.test_real_set[:len(occ_image_db. test_real_set) // 2] occ_val_set = LineModDatasetRealAug(occ_val_db, randomCropping, cfg.OCCLUSION_LINEMOD, vote_type, augment=False, cfg=train_cfg['aug_cfg'], use_motion=motion_model) occ_val_sampler = SequentialSampler(occ_val_set) occ_val_batch_sampler = ImageSizeBatchSampler( occ_val_sampler, train_cfg['test_batch_size'], False, cfg=train_cfg['aug_cfg']) occ_val_loader = DataLoader(occ_val_set, batch_sampler=occ_val_batch_sampler, num_workers=16) add_list_list = [] p_inc_list = [] p_dec_v_list = [] p_dec_q_list = [] largest_a_list = [] smallest_v_list = [] smallest_q_list = [] first_a_list = [] first_v_list = [] epoch_count = 0 for epoch in range(begin_epoch, train_cfg['epoch_num']): adjust_learning_rate(optimizer, epoch, train_cfg['lr_decay_rate'], train_cfg['lr_decay_epoch']) for param_group in optimizer.param_groups: lr = param_group['lr'] train(net, PVNet, optimizer, train_loader, epoch) # print('evaluate with test_loader') # _,_,_,_,_,_,_,_,_ = val(net, PVNet, test_loader, epoch, lr, writer, use_motion=motion_model) # print('evaluate with train_loader') # _,_,_,_,_,_,_,_,_ = val(net, PVNet, train_loader, epoch, lr, writer, use_motion=motion_model) print('evaluate with val_loader') add_list, first_a, first_v, largest_a, smallest_v, smallest_q, p_inc_add, p_dec_v, p_dec_q = val( net, PVNet, val_loader, epoch, lr, writer, use_motion=motion_model) if (train_cfg['eval_epoch'] and epoch % train_cfg['eval_inter'] == 0 and epoch >= train_cfg['eval_epoch_begin']) or args.test_model: if epoch >= 30: # add_list_list.append(add_list) # first_a_list.append(first_a) first_v_list.append(first_v) # p_inc_list.append(p_inc_add) p_dec_v_list.append(p_dec_v) p_dec_q_list.append(p_dec_q) # largest_a_list.append(largest_a) smallest_v_list.append(smallest_v) smallest_q_list.append(smallest_q) # if args.linemod_cls in cfg.occ_linemod_cls_names: # val(net, PVNet, occ_val_loader, epoch, lr, writer, 'occ_val',use_motion=motion_model) save_model(net.module.imNet, net.module.estNet, optimizer, epoch, model_dir) epoch_count += 1 print(train_cfg['exp_name']) # print('PVNet ADD. mean: {} +/- {}, max: {}'.format(np.mean(first_a_list),np.std(first_a_list),np.max(first_a_list))) # print('PVNet X-X^. mean: {} +/- {}, max: {}'.format(np.mean(first_v_list),np.std(first_v_list),np.max(first_v_list))) # print('ADD. mean: {} +/- {}, max: {}: '.format(np.mean(largest_a_list),np.std(largest_a_list),np.max(largest_a_list))) # print('ADD perc increase. mean: {} +/- {}, max: {}'.format(np.mean(p_inc_list),np.std(p_inc_list),np.max(p_inc_list))) print('X-X^ perc decrease. mean: {} +/- {}, max: {}'.format( np.mean(p_dec_v_list), np.std(p_dec_v_list), np.max(p_dec_v_list))) print('q-q^ perc decrease. mean: {} +/- {}, max: {}'.format( np.mean(p_dec_q_list), np.std(p_dec_q_list), np.max(p_dec_q_list)))