def initialize_loader(use_gpu): dataset_reid=initialize_single_batch() pin_memory = True if use_gpu else False transform_train,transform_test=set_transform() train_loader_reid = get_trainloader_resnetAttr(dataset_reid, transform_train, pin_memory) queryloader = DataLoader(ImageDataset(dataset_reid.query, transform=transform_test,arch=args.arch), batch_size=args.test_batch, shuffle=False, num_workers=args.workers,pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader(ImageDataset(dataset_reid.gallery, transform=transform_test,arch=args.arch), batch_size=args.test_batch, shuffle=False, num_workers=args.workers,pin_memory=pin_memory, drop_last=False, ) return [dataset_reid, train_loader_reid, queryloader, galleryloader]
def get_trainloader_resnetAttr(dataset_reid,transform_train,pin_memory): if args.dataset=='market1501' or args.dataset=='dukemtmcreid' or args.dataset=='pa100K': trainloader_reid = DataLoader( ImageDataset(dataset_reid.train,transform=transform_train,arch=args.arch), sampler=RandomIdentitySampler(dataset_reid.train, args.train_batch, args.num_instances), batch_size=args.train_batch, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) return trainloader_reid
def main(): torch.manual_seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = data_manager.init_imgreid_dataset( root=args.root, name=args.dataset, split_id=args.split_id, cuhk03_labeled=args.cuhk03_labeled, cuhk03_classic_split=args.cuhk03_classic_split, ) transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False trainloader = DataLoader( ImageDataset(dataset.train, transform=transform_train), batch_size=args.train_batch, shuffle=True, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) queryloader = DataLoader( ImageDataset(dataset.query, transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( ImageDataset(dataset.gallery, transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent'}, use_gpu=use_gpu) print("Model size: {:.3f} M".format(count_num_param(model))) if args.label_smooth: criterion = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) else: criterion = nn.CrossEntropyLoss() optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if args.fixbase_epoch > 0: if hasattr(model, 'classifier') and isinstance(model.classifier, nn.Module): optimizer_tmp = init_optim(args.optim, model.classifier.parameters(), args.fixbase_lr, args.weight_decay) else: print( "Warn: model has no attribute 'classifier' and fixbase_epoch is reset to 0" ) args.fixbase_epoch = 0 if args.load_weights and check_isfile(args.load_weights): # load pretrained weights but ignore layers that don't match in size checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format(args.load_weights)) if args.resume and check_isfile(args.resume): checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] + 1 best_rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch, best_rank1)) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") distmat = test(model, queryloader, galleryloader, use_gpu, return_distmat=True) if args.vis_ranked_res: visualize_ranked_results( distmat, dataset, save_dir=osp.join(args.save_dir, 'ranked_results'), topk=20, ) return start_time = time.time() train_time = 0 best_epoch = args.start_epoch print("==> Start training") if args.fixbase_epoch > 0: print( "Train classifier for {} epochs while keeping base network frozen". format(args.fixbase_epoch)) for epoch in range(args.fixbase_epoch): start_train_time = time.time() train(epoch, model, criterion, optimizer_tmp, trainloader, use_gpu, freeze_bn=True) train_time += round(time.time() - start_train_time) del optimizer_tmp print("Now open all layers for training") for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch + 1) > args.start_eval and args.eval_step > 0 and ( epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch: print("==> Test") rank1 = test(model, queryloader, galleryloader, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format( best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = DatasetManager(dataset_dir=args.dataset, root=args.root) transform_train = T.Compose_Keypt([ T.Random2DTranslation_Keypt((args.width, args.height)), T.RandomHorizontalFlip_Keypt(), T.ToTensor_Keypt(), T.Normalize_Keypt(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose_Keypt([ T.Resize_Keypt((args.width, args.height)), T.ToTensor_Keypt(), T.Normalize_Keypt(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False trainloader = DataLoader( ImageDataset(dataset.train, keyptaware=args.keyptaware, heatmapaware=args.heatmapaware, segmentaware=args.segmentaware, transform=transform_train, imagesize=(args.width, args.height)), sampler=RandomIdentitySampler(dataset.train, args.train_batch, args.num_instances), batch_size=args.train_batch, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) queryloader = DataLoader( ImageDataset(dataset.query, keyptaware=args.keyptaware, heatmapaware=args.heatmapaware, segmentaware=args.segmentaware, transform=transform_test, imagesize=(args.width, args.height)), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( ImageDataset(dataset.gallery, keyptaware=args.keyptaware, heatmapaware=args.heatmapaware, segmentaware=args.segmentaware, transform=transform_test, imagesize=(args.width, args.height)), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_vids=dataset.num_train_vids, num_vcolors=dataset.num_train_vcolors, num_vtypes=dataset.num_train_vtypes, keyptaware=args.keyptaware, heatmapaware=args.heatmapaware, segmentaware=args.segmentaware, multitask=args.multitask) print("Model size: {:.3f} M".format(count_num_param(model))) if args.label_smooth: criterion_xent_vid = CrossEntropyLabelSmooth( num_classes=dataset.num_train_vids, use_gpu=use_gpu) criterion_xent_vcolor = CrossEntropyLabelSmooth( num_classes=dataset.num_train_vcolors, use_gpu=use_gpu) criterion_xent_vtype = CrossEntropyLabelSmooth( num_classes=dataset.num_train_vtypes, use_gpu=use_gpu) else: criterion_xent_vid = nn.CrossEntropyLoss() criterion_xent_vcolor = nn.CrossEntropyLoss() criterion_xent_vtype = nn.CrossEntropyLoss() criterion_htri = TripletLoss(margin=args.margin) optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if args.load_weights: # load pretrained weights but ignore layers that don't match in size if check_isfile(args.load_weights): checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format( args.load_weights)) if args.resume: if check_isfile(args.resume): checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format( args.start_epoch, rank1)) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") distmat = test(model, args.keyptaware, args.multitask, queryloader, galleryloader, use_gpu, dataset.vcolor2label, dataset.vtype2label, return_distmat=True) if args.vis_ranked_res: visualize_ranked_results( distmat, dataset, save_dir=osp.join(args.save_dir, 'ranked_results'), topk=100, ) return start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 print("==> Start training") for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, args.keyptaware, args.multitask, criterion_xent_vid, criterion_xent_vcolor, criterion_xent_vtype, criterion_htri, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if ((epoch + 1) > args.start_eval and args.eval_step > 0 and (epoch + 1) % args.eval_step == 0) or (epoch + 1) == args.max_epoch: print("==> Test") rank1 = test(model, args.keyptaware, args.multitask, queryloader, galleryloader, use_gpu, dataset.vcolor2label, dataset.vtype2label) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() # rank1 = 1 # is_best = True save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("==> Best Rank-1 {:.2%}, achieved at epoch {}".format( best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = data_manager.init_imgreid_dataset( root=args.root, name=args.dataset, split_id=args.split_id, cuhk03_labeled=args.cuhk03_labeled, cuhk03_classic_split=args.cuhk03_classic_split, ) #cuhk03_labeled: detected,labeled transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False #pdb.set_trace() trainloader = DataLoader( ImageDataset(dataset.train, transform=transform_train), sampler=RandomIdentitySampler(dataset.train, args.train_batch, args.num_instances), batch_size=args.train_batch, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) queryloader = DataLoader( ImageDataset(dataset.query, transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( ImageDataset(dataset.gallery, transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, batchsize=args.test_batch, loss={'xent', 'wcont', 'htri'}) print("Model size: {:.3f} M".format(count_num_param(model))) criterion_xent = nn.CrossEntropyLoss() criterion_htri = TripletLoss(margin=args.margin) criterion_KA = KALoss(margin=args.margin, same_margin=args.same_margin, use_auto_samemargin=args.use_auto_samemargin) cirterion_lifted = LiftedLoss(margin=args.margin) cirterion_batri = BA_TripletLoss(margin=args.margin) if args.use_auto_samemargin == True: G_params = [{ 'params': model.parameters(), 'lr': args.lr }, { 'params': criterion_KA.auto_samemargin, 'lr': args.lr }] else: G_params = [para for _, para in model.named_parameters()] optimizer = init_optim(args.optim, G_params, args.lr, args.weight_decay) if args.load_weights: # load pretrained weights but ignore layers that don't match in size if check_isfile(args.load_weights): checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format( args.load_weights)) if args.resume: if check_isfile(args.resume): checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format( args.start_epoch, rank1)) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") distmat = test(model, queryloader, galleryloader, use_gpu, return_distmat=True) if args.vis_ranked_res: visualize_ranked_results( distmat, dataset, save_dir=osp.join(args.save_dir, 'ranked_results'), topk=20, ) return start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 print("==> Start training") for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() adjust_learning_rate(optimizer, epoch) train(epoch, model, cirterion_batri, cirterion_lifted, criterion_xent, criterion_htri, criterion_KA, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) if (epoch + 1) > args.start_eval and args.eval_step > 0 and ( epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch: rank1 = 0 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() print("==> Test") sys.stdout.flush() rank1 = test(model, queryloader, galleryloader, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("model saved") print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format( best_rank1, best_epoch)) sys.stdout.flush() elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time)) sys.stdout.flush()
def main(): global args, best_rank1 torch.manual_seed(args.seed) if not args.use_avai_gpus: os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = data_manager.init_vidreid_dataset(root=args.root, name=args.dataset) transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False # decompose tracklets into images for image-based training new_train = [] for img_paths, pid, camid in dataset.train: for img_path in img_paths: new_train.append((img_path, pid, camid)) trainloader = DataLoader( ImageDataset(new_train, transform=transform_train), sampler=RandomIdentitySampler(new_train, args.train_batch, args.num_instances), batch_size=args.train_batch, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) queryloader = DataLoader( VideoDataset(dataset.query, seq_len=args.seq_len, sample='evenly', transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( VideoDataset(dataset.gallery, seq_len=args.seq_len, sample='evenly', transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}) print("Model size: {:.3f} M".format(count_num_param(model))) if args.label_smooth: criterion_xent = CrossEntropyLabelSmooth( num_classes=dataset.num_train_pids, use_gpu=use_gpu) else: criterion_xent = nn.CrossEntropyLoss() criterion_htri = TripletLoss(margin=args.margin) optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if args.load_weights and check_isfile(args.load_weights): # load pretrained weights but ignore layers that don't match in size checkpoint = torch.load(args.load_weights) pretrain_dict = checkpoint['state_dict'] model_dict = model.state_dict() pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict and model_dict[k].size() == v.size() } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) print("Loaded pretrained weights from '{}'".format(args.load_weights)) if args.resume and check_isfile(args.resume): checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) args.start_epoch = checkpoint['epoch'] + 1 best_rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(args.start_epoch, best_rank1)) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") distmat = test(model, queryloader, galleryloader, args.pool, use_gpu, return_distmat=True) if args.visualize_ranks: visualize_ranked_results( distmat, dataset, save_dir=osp.join(args.save_dir, 'ranked_results'), topk=20, ) return start_time = time.time() train_time = 0 best_epoch = args.start_epoch print("==> Start training") for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch + 1) > args.start_eval and args.eval_step > 0 and ( epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch: print("==> Test") rank1 = test(model, queryloader, galleryloader, args.pool, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': rank1, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format( best_rank1, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
def main(): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices use_gpu = torch.cuda.is_available() if args.use_cpu: use_gpu = False if not args.evaluate: sys.stdout = Logger(osp.join(args.save_dir, 'log_train.txt')) else: sys.stdout = Logger(osp.join(args.save_dir, 'log_test.txt')) print("==========\nArgs:{}\n==========".format(args)) if use_gpu: print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) else: print("Currently using CPU (GPU is highly recommended)") print("Initializing dataset {}".format(args.dataset)) dataset = data_manager.init_imgreid_dataset(root=args.root, name=args.dataset, split_id=args.split_id) transform_train = T.Compose([ T.Random2DTranslation(args.height, args.width), T.RandomHorizontalFlip(), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) transform_test = T.Compose([ T.Resize((args.height, args.width)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True if use_gpu else False trainloader = DataLoader( ImageDataset(dataset.train, transform=transform_train), batch_size=args.train_batch, shuffle=True, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) testloader = DataLoader( ImageDataset(dataset.test, transform=transform_test), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing model: {}".format(args.arch)) model = models.init_model(name=args.arch, loss={'xent'}, use_gpu=use_gpu) print("Model size: {:.3f} M".format(count_num_param(model))) gender_criterion_xent = nn.CrossEntropyLoss() staff_criterion_xent = nn.CrossEntropyLoss() customer_criterion_xent = nn.CrossEntropyLoss() stand_criterion_xent = nn.CrossEntropyLoss() sit_criterion_xent = nn.CrossEntropyLoss() phone_criterion_xent = nn.CrossEntropyLoss() optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) if use_gpu: model = nn.DataParallel(model).cuda() start_time = time.time() train_time = 0 best_score = 0 best_epoch = args.start_epoch print("==> Start training") ################################### 修改到这里,把train 和 test改一下就好 for epoch in range(args.start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, gender_criterion_xent, staff_criterion_xent, customer_criterion_xent, \ stand_criterion_xent, sit_criterion_xent, phone_criterion_xent, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) scheduler.step() if (epoch + 1) > args.start_eval and args.eval_step > 0 and ( epoch + 1) % args.eval_step == 0 or (epoch + 1) == args.max_epoch: print("==> Test") gender_accurary, staff_accurary, customer_accurary, stand_accurary, sit_accurary, phone_accurary = test( model, testloader, use_gpu) Score = (gender_accurary + staff_accurary + customer_accurary + stand_accurary + sit_accurary + phone_accurary) * 100 is_best = Score > best_score if is_best: best_score = Score best_gender_acc = gender_accurary best_staff_acc = staff_accurary best_customer_acc = customer_accurary best_stand_acc = stand_accurary best_sit_acc = sit_accurary best_phone_acc = phone_accurary best_epoch = epoch + 1 if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() save_checkpoint( { 'state_dict': state_dict, 'rank1': Score, 'epoch': epoch, }, is_best, osp.join(args.save_dir, 'checkpoint_ep' + str(epoch + 1) + '.pth.tar')) print( "==> Best best_score {} |Gender_acc {}\t Staff_acc {}\t Customer_acc {}\t Stand_acc {}\t Sit_acc {}\t Phone_acc {}|achieved at epoch {}" .format(best_score, best_gender_acc, best_staff_acc, best_customer_acc, best_stand_acc, best_sit_acc, best_phone_acc, best_epoch)) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) train_time = str(datetime.timedelta(seconds=train_time)) print( "Finished. Total elapsed time (h:m:s): {}. Training time (h:m:s): {}.". format(elapsed, train_time))
print("Currently using GPU {}".format(args.gpu_devices)) cudnn.benchmark = True torch.cuda.manual_seed_all(args.seed) device = 'cuda' else: device = 'cpu' print("Currently using CPU, however, GPU is highly recommended") dataset = init_imgreid_dataset(args.source_names[0], root=args.root) transform_test_flip = build_transforms(args.height, args.width, is_train=False, data_augment=args.data_augment, flip=True) val_set = ImageDataset(dataset.query + dataset.gallery, transform_test_flip) val_loader = DataLoader(val_set, batch_size=args.test_batch_size, shuffle=False, num_workers=args.workers, collate_fn=val_collate_fn) print(val_loader) model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent'}, use_gpu=use_gpu, args=vars(args)) if args.load_weights and check_isfile(args.load_weights):
def __init__( self, use_gpu, source_names, target_names, root, split_id=0, height=256, width=128, train_batch_size=32, test_batch_size=100, workers=4, train_sampler='', data_augment='none', num_instances=4, # number of instances per identity (for RandomIdentitySampler) cuhk03_labeled=False, # use cuhk03's labeled or detected images cuhk03_classic_split=False # use cuhk03's classic split or 767/700 split ): super(ImageDataManager, self).__init__() from torchreid.dataset_loader import ImageDataset from torchreid.datasets import init_imgreid_dataset from torchreid.transforms import build_transforms from torch.utils.data import DataLoader from torchreid.samplers import RandomIdentitySampler self.use_gpu = use_gpu self.source_names = source_names self.target_names = target_names self.root = root self.split_id = split_id self.height = height self.width = width self.train_batch_size = train_batch_size self.test_batch_size = test_batch_size self.workers = workers self.train_sampler = train_sampler self.num_instances = num_instances self.cuhk03_labeled = cuhk03_labeled self.cuhk03_classic_split = cuhk03_classic_split self.pin_memory = True if self.use_gpu else False # Build train and test transform functions transform_train = build_transforms(self.height, self.width, is_train=True, data_augment=data_augment) transform_test = build_transforms(self.height, self.width, is_train=False, data_augment=data_augment) # transform_test_flip = build_transforms(self.height, self.width, is_train=False, data_augment=data_augment, flip=True) print("=> Initializing TRAIN (source) datasets") self.train = [] self._num_train_pids = 0 self._num_train_cams = 0 for name in self.source_names: dataset = init_imgreid_dataset( root=self.root, name=name, split_id=self.split_id, cuhk03_labeled=self.cuhk03_labeled, cuhk03_classic_split=self.cuhk03_classic_split) for img_path, pid, camid in dataset.train: pid += self._num_train_pids camid += self._num_train_cams self.train.append((img_path, pid, camid)) self._num_train_pids += dataset.num_train_pids self._num_train_cams += dataset.num_train_cams if self.train_sampler == 'RandomIdentitySampler': print('!!! Using RandomIdentitySampler !!!') self.trainloader = DataLoader( ImageDataset(self.train, transform=transform_train), sampler=RandomIdentitySampler(self.train, self.train_batch_size, self.num_instances), batch_size=self.train_batch_size, shuffle=False, num_workers=self.workers, pin_memory=self.pin_memory, drop_last=True) else: self.trainloader = DataLoader(ImageDataset( self.train, transform=transform_train), batch_size=self.train_batch_size, shuffle=True, num_workers=self.workers, pin_memory=self.pin_memory, drop_last=True) print("=> Initializing TEST (target) datasets") self.testloader_dict = { name: { 'query': None, 'gallery': None } for name in self.target_names } self.testdataset_dict = { name: { 'query': None, 'gallery': None } for name in self.target_names } for name in self.target_names: dataset = init_imgreid_dataset( root=self.root, name=name, split_id=self.split_id, cuhk03_labeled=self.cuhk03_labeled, cuhk03_classic_split=self.cuhk03_classic_split) self.testloader_dict[name]['new_vid_old_cid_query'] = DataLoader( ImageDataset(dataset.new_vid_old_cid_query, transform=transform_test), batch_size=self.test_batch_size, shuffle=False, num_workers=self.workers, pin_memory=self.pin_memory, drop_last=False) self.testloader_dict[name]['new_vid_old_cid_val'] = DataLoader( ImageDataset(dataset.new_vid_old_cid_val, transform=transform_test), batch_size=self.test_batch_size, shuffle=False, num_workers=self.workers, pin_memory=self.pin_memory, drop_last=False) self.testloader_dict[name]['new_vid_new_cid_query'] = DataLoader( ImageDataset(dataset.new_vid_new_cid_query, transform=transform_test), batch_size=self.test_batch_size, shuffle=False, num_workers=self.workers, pin_memory=self.pin_memory, drop_last=False) self.testloader_dict[name]['new_vid_new_cid_val'] = DataLoader( ImageDataset(dataset.new_vid_new_cid_val, transform=transform_test), batch_size=self.test_batch_size, shuffle=False, num_workers=self.workers, pin_memory=self.pin_memory, drop_last=False) self.testloader_dict[name]['train_gallery'] = DataLoader( ImageDataset(dataset.train_gallery, transform=transform_test), batch_size=self.test_batch_size, shuffle=False, num_workers=self.workers, pin_memory=self.pin_memory, drop_last=False) # self.testdataset_dict[name]['query'] = dataset.query # self.testdataset_dict[name]['gallery'] = dataset.gallery print("\n") print(" **************** Summary ****************") print(" train names : {}".format(self.source_names)) print(" # train datasets : {}".format(len(self.source_names))) print(" # train ids : {}".format(self._num_train_pids)) print(" # train images : {}".format(len(self.train))) print(" # train cameras : {}".format(self._num_train_cams)) print(" test names : {}".format(self.target_names)) print(" *****************************************") print("\n")