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_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 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), batch_size=args.train_batch, shuffle=True, 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'}) print("Model size: {:.3f} M".format(count_num_param(model))) criterion = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) start_epoch = args.start_epoch if args.fixbase_epoch > 0: optimizer_tmp = init_optim(args.optim, model.classifier.parameters(), args.fixbase_lr, args.weight_decay) if args.resume: print("Loading checkpoint from '{}'".format(args.resume)) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") test(model, queryloader, galleryloader, args.pool, use_gpu) return start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 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(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, 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_img_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.Resize((args.height, args.width)), T.RandomHorizontalFlip(p=0.5), T.Pad(10), T.RandomCrop([args.height, args.width]), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), torchvision.transforms.RandomErasing(p=0.5, scale=(0.02, 0.4), ratio=(0.3, 3.33), value=(0.4914, 0.4822, 0.4465)) # T.RandomErasing(probability=0.5, sh=0.4, mean=(0.4914, 0.4822, 0.4465)), ]) 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={'cent'}) print("Model size: {:.5f}M".format( sum(p.numel() for p in model.parameters()) / 1000000.0)) criterion_xent = CrossEntropyLabelSmooth( num_classes=dataset.num_train_pids, use_gpu=use_gpu) criterion_cent = CenterLoss(num_classes=dataset.num_train_pids, feat_dim=model.feat_dim, use_gpu=use_gpu) optimizer_model = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) optimizer_cent = torch.optim.SGD(criterion_cent.parameters(), lr=args.lr_cent) #only the optimizer_model use learning rate schedule # if args.stepsize > 0: # scheduler = lr_scheduler.StepLR(optimizer_model, step_size=args.stepsize, gamma=args.gamma) '''------Modify lr_schedule here------''' current_schedule = init_lr_schedule(schedule=args.schedule, warm_up_epoch=args.warm_up_epoch, half_cos_period=args.half_cos_period, lr_milestone=args.lr_milestone, gamma=args.gamma, stepsize=args.stepsize) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer_model, lr_lambda=current_schedule) '''------Please refer to the args.xxx for details of hyperparams------''' # embed() start_epoch = args.start_epoch if args.resume: print("Loading checkpoint from '{}'".format(args.resume)) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") test(model, queryloader, galleryloader, use_gpu) return start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 print("==> Start training") for epoch in range(start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, criterion_xent, criterion_cent, optimizer_model, optimizer_cent, trainloader, use_gpu) train_time += round(time.time() - start_train_time) if args.schedule: 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')) # tensorboardX # writer = SummaryWriter(log_dir=osp.join(args.save_dir,'summary')) 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_img_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]), ]) if args.random_erasing: 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]), RandomErasing(probability=args.probability, mean=[0.0, 0.0, 0.0]), ]) 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 if args.loss == 'xent,htri': trainloader = DataLoader( ImageDataset(dataset.train, transform=transform_train), sampler=RandomIdentitySampler(dataset.train, num_instances=args.num_instances), batch_size=args.train_batch, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) elif args.loss == 'xent': 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=args.loss) print("Model size: {:.5f}M".format(sum(p.numel() for p in model.parameters())/1000000.0)) criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) criterion_htri = TripletLoss(margin=args.margin) optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) if args.stepsize > 0: if not args.warmup: scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma) start_epoch = args.start_epoch if args.resume: print("Loading checkpoint from '{}'".format(args.resume)) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") test(model, queryloader, galleryloader, use_gpu) return def adjust_lr(optimizer, ep): if ep < 20: lr = 1e-4 * (ep + 1) / 2 elif ep < 80: #lr = 1e-3 * len(args.gpu_devices) lr = 1e-3 elif ep < 180: #lr = 1e-4 * len(args.gpu_devices) lr = 1e-4 elif ep < 300: #lr = 1e-5 * len(args.gpu_devices) lr = 1e-5 elif ep < 320: #lr = 1e-5 * 0.1 ** ((ep - 320) / 80) * len(args.gpu_devices) lr = 1e-5 * 0.1 ** ((ep - 320) / 80) elif ep < 400: lr = 1e-6 elif ep < 480: #lr = 1e-4 * len(args.gpu_devices) lr = 1e-4 else: #lr = 1e-5 * len(args.gpu_devices) lr = 1e-5 for p in optimizer.param_groups: p['lr'] = lr length = len(trainloader) start_time = time.time() train_time = 0 best_rank1 = -np.inf best_epoch = 0 #best_rerank1 = -np.inf #best_rerankepoch = 0 print("==> Start training") for epoch in range(start_epoch, args.max_epoch): start_train_time = time.time() if args.stepsize > 0: if args.warmup: adjust_lr(optimizer, epoch + 1) else: scheduler.step() train(epoch, model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu=use_gpu, summary=None, length=length) 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: print("==> Test") rank1 = test(epoch, model, queryloader, galleryloader, use_gpu=True, summary=None) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 best_epoch = epoch + 1 ####### Best Rerank #is_rerankbest = rerank1 > best_rerank1 #if is_rerankbest: # best_rerank1 = rerank1 # best_rerankepoch = 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')) writer.close() print("==> Best Rank-1 {:.1%}, achieved at epoch {}".format(best_rank1, best_epoch)) #print("==> Best Rerank-1 {:.1%}, achieved at epoch {}".format(best_rerank1, best_rerankepoch)) 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))
return coloured_mask img_paths = glob.glob("../main20170707/org_imgs/*.png") # img_paths = glob.glob("../data/tool/org_imgs/*.png") # dataset = Dataset(img_paths, annotation, is_train=False) model.load_state_dict(torch.load("model.pth", map_location=device)) model.eval() confidence = 0.5 # for idx in tqdm.tqdm(range(dataset.__len__())): for idx in tqdm(range(len(img_paths))): # Prediction # img, _ = dataset.__getitem__(idx) img_path = img_paths[idx] img = Image.open(img_path) transform = T.Compose([T.ToTensor()]) img = torchvision.transforms.functional.to_tensor(img) pred = model([img.to(device)]) pred_score = list(pred[0]['scores'].detach().cpu().numpy()) pred_t = [ pred_score.index(x) for x in pred_score if x > confidence ] masks = (pred[0]['masks'] > 0.5).squeeze().detach().cpu().numpy() if masks.ndim == 2: masks = masks.reshape([1, masks.shape[0], masks.shape[1]]) pred_class = [ CLASS_NAMES[i] for i in list(pred[0]['labels'].cpu().numpy()) ] pred_boxes = [[ (i[0], i[1]), (i[2], i[3])
from PIL import Image import torch import torch.nn as nn import torch.nn.functional as F import os from torch.autograd import Variable import transforms as transforms from skimage import io from skimage.transform import resize from models import * import cv2 as cv import dlib import time transform_test = transforms.Compose([ transforms.ToTensor(), ]) def rect_to_bb(rect): x = rect.left() y = rect.top() w = rect.right() - x h = rect.bottom() - y # return a tuple of (x, y, w, h) return x, y, w, h def apply_offsets(face_coordinates, offsets): x, y, width, height = face_coordinates
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_dataset(root=args.root, name=args.dataset, cls_sample=args.cls_sample) # print(dataset.train) # print(1) # 解释器:创建一个transform处理图像数据的设置 # T.Random2DTranslation:随机裁剪 # T.RandomHorizontalFlip: 给定概率进行随机水平翻转 # T.ToTensor: 将PIL或numpy向量[0,255]=>tensor[0.0,1.0] # T.Normalize:用均值和标准偏差标准化张量图像,mean[ , , ]三个参数代表三通道 transform_train = T.Compose([ # T.RandomCrop(224), 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(256), # T.CenterCrop(224), 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 m = dataset.train print(1) # Dataloader 提供队列和线程 # ImageDataset:return data =>img, pid, camid # RandomIdentitySampler:定义从数据集中抽取样本的策略 # num_workers: 子进程数 # print(dataset.train) trainloader = DataLoader( AGE_Gender_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( AGE_Gender_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, num_classes=dataset.train_num_class, loss={'xent'}) print("Model size: {:.5f}M".format( sum(p.numel() for p in model.parameters()) / 1000000.0)) # criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset.train_num_class, use_gpu=use_gpu) age_criterion_xent = nn.CrossEntropyLoss() gender_criterion_xent = nn.CrossEntropyLoss() # optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay) if args.stepsize > 0: scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma) start_epoch = args.start_epoch if args.resume: print("Loading checkpoint from '{}'".format(args.resume)) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") test(model, testloader, use_gpu) return start_time = time.time() train_time = 0 # best_rank1 = -np.inf best_score = 0 best_MAE = 0 best_gender_acc = 0 best_epoch = 0 print("==> Start training") for epoch in range(start_epoch, args.max_epoch): start_train_time = time.time() train(epoch, model, age_criterion_xent, gender_criterion_xent, optimizer, trainloader, use_gpu) train_time += round(time.time() - start_train_time) if args.stepsize > 0: scheduler.step() if args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or ( epoch + 1) == args.max_epoch: print("==> Test") MAE, Gender_acc = test(model, testloader, use_gpu) Score = Gender_acc * 100 - MAE is_best = Score > best_score if is_best: best_score = Score best_MAE = MAE best_gender_acc = Gender_acc 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-MAE) {} |Gender_acc {}\t MAE {}|achieved at epoch {}" .format(best_score, best_gender_acc, best_MAE, 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))
"""Keep reading image until succeed. This can avoid IOError incurred by heavy IO process.""" got_img = False if not osp.exists(img_path): raise IOError("{} does not exist".format(img_path)) while not got_img: try: img = Image.open(img_path).convert('RGB') got_img = True except IOError: print("IOError incurred when reading '{}'. Will redo. Don't worry. Just chill.".format(img_path)) pass return img transform_test = T.Compose([ T.Resize((args.height, args.width), interpolation=3), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) def train(epoch,model_edge, model, model_tradclass,weight_softmax, criterion, optimizer, trainloader, learning_rate, use_gpu): if not os.path.isdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_1"): os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_1") os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_2") os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_3") os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_4") os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_5") os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_6") os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_7") os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_8") os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_9") os.mkdir("/data4/lijunjie/tiered-imagenet-tools-master/tiered_imagenet/train_10")
def testseq(dataset_name, use_gpu): dataset_root = './video2img/track1_sct_img_test_big/' dataset = Graph_data_manager.AICityTrack2(root=dataset_root) width = 224 height = 224 transform_train = T.Compose([ T.Random2DTranslation(height, 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((height, 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 seq_len = 4 num_instance = 4 train_batch = 32 test_batch = 1 queryloader = DataLoader( VideoDataset(dataset.query, seq_len=seq_len, sample='dense', transform=transform_test), batch_size=test_batch, shuffle=False, num_workers=4, pin_memory=pin_memory, drop_last=False, ) arch = "resnet50ta" pretrained_model = "./log/track12_ta224_checkpoint_ep500.pth.tar" start_epoch = 0 print("Initializing model: {}".format(arch)) dataset.num_train_pids = 517 if arch == 'resnet503d': model = resnet3d.resnet50(num_classes=dataset.num_train_pids, sample_width=width, sample_height=height, sample_duration=seq_len) if not os.path.exists(pretrained_model): raise IOError( "Can't find pretrained model: {}".format(pretrained_model)) print("Loading checkpoint from '{}'".format(pretrained_model)) checkpoint = torch.load(pretrained_model) state_dict = {} for key in checkpoint['state_dict']: if 'fc' in key: continue state_dict[key.partition("module.") [2]] = checkpoint['state_dict'][key] model.load_state_dict(state_dict, strict=False) else: if not os.path.exists(pretrained_model): model = models.init_model(name=arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}) else: model = models.init_model(name=arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}) checkpoint = torch.load(pretrained_model) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] + 1 print("Loaded checkpoint from '{}'".format(pretrained_model)) print("- start_epoch: {}\n- rank1: {}".format( start_epoch, checkpoint['rank1'])) print("Model size: {:.5f}M".format( sum(p.numel() for p in model.parameters()) / 1000000.0)) criterion_xent = CrossEntropyLabelSmooth( num_classes=dataset.num_train_pids, use_gpu=use_gpu) criterion_htri = TripletLoss(margin=0.3) lr = 0.0003 gamma = 0.1 stepsize = 200 weight_decay = 5e-04 optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay) if stepsize > 0: scheduler = lr_scheduler.StepLR(optimizer, step_size=stepsize, gamma=gamma) start_epoch = start_epoch if use_gpu: model = nn.DataParallel(model).cuda() test(model, queryloader, 'avg', use_gpu, dataset, -1, meta_data_tab=None)
def main(parser_data): device = torch.device( parser_data.device if torch.cuda.is_available() else "cpu") print("Using {} device training.".format(device.type)) data_transform = {"val": transforms.Compose([transforms.ToTensor()])} # read class_indict label_json_path = './pascal_voc_classes.json' assert os.path.exists( label_json_path), "json file {} dose not exist.".format( label_json_path) json_file = open(label_json_path, 'r') class_dict = json.load(json_file) category_index = {v: k for k, v in class_dict.items()} VOC_root = parser_data.data_path # check voc root if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False: raise FileNotFoundError( "VOCdevkit dose not in path:'{}'.".format(VOC_root)) # 注意这里的collate_fn是自定义的,因为读取的数据包括image和targets,不能直接使用默认的方法合成batch batch_size = parser_data.batch_size nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers print('Using %g dataloader workers' % nw) # load validation data set val_data_set = VOC2007DataSet(VOC_root, data_transform["val"], "val.txt") val_data_set_loader = torch.utils.data.DataLoader( val_data_set, batch_size=batch_size, shuffle=False, num_workers=nw, collate_fn=val_data_set.collate_fn) # create model num_classes equal background + 20 classes backbone = resnet50_fpn_backbone() model = FasterRCNN(backbone=backbone, num_classes=parser_data.num_classes + 1) # 载入你自己训练好的模型权重 weights_path = parser_data.weights assert os.path.exists(weights_path), "not found {} file.".format( weights_path) weights_dict = torch.load(weights_path, map_location=device) model.load_state_dict(weights_dict['model']) # print(model) model.to(device) # evaluate on the test dataset coco = get_coco_api_from_dataset(val_data_set) iou_types = ["bbox"] coco_evaluator = CocoEvaluator(coco, iou_types) cpu_device = torch.device("cpu") model.eval() with torch.no_grad(): for image, targets in tqdm(val_data_set_loader, desc="validation..."): # 将图片传入指定设备device image = list(img.to(device) for img in image) # inference outputs = model(image) outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs] res = { target["image_id"].item(): output for target, output in zip(targets, outputs) } coco_evaluator.update(res) coco_evaluator.synchronize_between_processes() # accumulate predictions from all images coco_evaluator.accumulate() coco_evaluator.summarize() coco_eval = coco_evaluator.coco_eval["bbox"] # calculate COCO info for all classes coco_stats, print_coco = summarize(coco_eval) # calculate voc info for every classes(IoU=0.5) voc_map_info_list = [] for i in range(len(category_index)): stats, _ = summarize(coco_eval, catId=i) voc_map_info_list.append(" {:15}: {}".format(category_index[i + 1], stats[1])) print_voc = "\n".join(voc_map_info_list) print(print_voc) # 将验证结果保存至txt文件中 with open("record_mAP.txt", "w") as f: record_lines = [ "COCO results:", print_coco, "", "mAP(IoU=0.5) for each category:", print_voc ] f.write("\n".join(record_lines))
def main(**kwargs): training_path = kwargs.get('training_data_path') checkpoint_path = kwargs.get('checkpoint_path') tensorboard_log_dir = kwargs.get('tensorboard_log_dir') if not os.path.isdir(checkpoint_path): os.mkdir(checkpoint_path) backbone_name = kwargs.get('backbone') criterion_name = kwargs.get('criterion').upper() optimizer_name = kwargs.get('optimizer').upper() scheduler = kwargs.get('scheduler',None) pretrained = kwargs.get('pretrained') num_classes = kwargs.get('num_classes') device = kwargs.get('device') batch_size = kwargs.get('batch_size') epochs = kwargs.get('epochs') hyperparameters = kwargs.get('hyperparameters',{}) augmentations = kwargs.get('augmentations',{}) verbose = kwargs.get('verbose') train_split = kwargs.get('train_split') nfolds = kwargs.get('nfolds') val_splits = [(1-train_split) / nfolds] * nfolds resume = kwargs.get('resume') only_weights = kwargs.get('only_weights') seed = hyperparameters.get('seed') random_jitter = augmentations.get('jitter',{}) random_horizontal_flip = augmentations.get('horizontal_flip', 0.5) random_rotation = augmentations.get('rotation', 20) writer = SummaryWriter(log_dir=tensorboard_log_dir, flush_secs=20) if seed: seed_everything(seed) # TODO calculate mean and std mean = hyperparameters.get('mean',0) std = hyperparameters.get('std',1) splits = [train_split]+val_splits assert sum(splits) <= 1,"given splits must be lower or equal than 1" original_img_size = 96 criterion = get_criterion(criterion_name) backbone = get_backbone(backbone_name, pretrained=pretrained) model = FacialKeypointsDetector(backbone, criterion=criterion, device=device, num_classes=num_classes) optimizer = get_optimizer(optimizer_name, model.parameters(), kwargs=hyperparameters.get('optimizer',{})) scaler = GradScaler() val_transforms = None val_target_transform = TargetTransform(original_img_size) train_transform = train_target_transform = None train_transforms = transforms.TrainTransforms(model.get_input_size(), original_img_size, mean=mean, std=std, brightness=random_jitter.get('brightness'), contrast=random_jitter.get('contrast'), saturation=random_jitter.get('saturation'), hue=random_jitter.get('hue'), rotation_degree=random_rotation, hflip=random_horizontal_flip) val_transform = transforms.Compose([ transforms.ToPILImage(), transforms.Resize(model.get_input_size()), transforms.ToTensor(), transforms.Normalize(mean,std)]) train_dataset,*val_datasets = get_training_datasets(root_path=training_path, train_transforms=(train_transforms,train_transform,train_target_transform), val_transforms=(val_transforms,val_transform,val_target_transform), split_ratios=splits) val_dls = [] train_dl = torch.utils.data.DataLoader(train_dataset, num_workers=4, batch_size=batch_size, pin_memory=True, collate_fn=custom_collate_fn, shuffle=True) for val_ds in val_datasets: val_dls.append( torch.utils.data.DataLoader( val_ds, batch_size=batch_size, num_workers=2) ) current_epoch = 0 best_loss = math.inf if resume: print(Fore.CYAN, f"loading checkpoint from {checkpoint_path}",Style.RESET_ALL) best_loss,current_epoch = load_checkpoint(model, optimizer, scheduler=scheduler, save_path=checkpoint_path, suffix='last', only_weights=only_weights) try: for epoch in range(current_epoch,epochs): training_loop(train_dl, model, epoch, epochs, optimizer, writer,scaler, scheduler=scheduler, verbose=verbose) val_losses = [] for i,val_dl in enumerate(val_dls): val_loss = validation_loop(val_dl, model) val_losses.append(val_loss) print(Fore.LIGHTBLUE_EX, f"validation [{i+1}] loss: {val_loss:.07f}",Style.RESET_ALL) writer.add_scalar(f'Loss/val_{i+1}', val_loss, epoch) mean_val_loss = sum(val_losses) / len(val_losses) print(Fore.LIGHTBLUE_EX, f"validation [mean] loss: {mean_val_loss:.07f}",Style.RESET_ALL) writer.add_scalar(f'Loss/val_mean', mean_val_loss, epoch) writer.flush() if mean_val_loss < best_loss: best_loss = mean_val_loss print(Fore.CYAN, "saving best checkpoint...",Style.RESET_ALL) save_checkpoint(model,optimizer,epoch,best_loss, scheduler=scheduler, suffix='best', save_path=checkpoint_path) print(Fore.CYAN, "saving last checkpoint...",Style.RESET_ALL) save_checkpoint(model,optimizer,epoch,best_loss, scheduler=scheduler, suffix='last', save_path=checkpoint_path) except KeyboardInterrupt: print(Fore.RED, "training interrupted with ctrl+c saving current state of the model",Style.RESET_ALL) save_checkpoint(model,optimizer,epoch,best_loss, scheduler=scheduler, suffix='last', save_path=checkpoint_path) writer.close()
def main(): global best_acc start_epoch = args.start_epoch # start from epoch 0 or last checkpoint epoch if not os.path.isdir(args.checkpoint): mkdir_p(args.checkpoint) # Data print('==> Preparing dataset %s' % args.dataset) transform_train = transforms.Compose([ transforms.RandomCrop(28, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.1307, ), (0.3081, )), transforms.RandomErasing(probability=args.p, sh=args.sh, r1=args.r1, mean=[0.4914]), ]) transform_test = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307, ), (0.3081, )), ]) if args.dataset == 'fashionmnist': dataloader = datasets.FashionMNIST num_classes = 10 trainset = dataloader(root='../../dataset', train=True, download=True, transform=transform_train) trainloader = data.DataLoader(trainset, batch_size=args.train_batch, shuffle=True, num_workers=args.workers) testset = dataloader(root='../../dataset', train=False, download=False, transform=transform_test) testloader = data.DataLoader(testset, batch_size=args.test_batch, shuffle=False, num_workers=args.workers) # Model print("==> creating model '{}'".format(args.arch)) if args.arch.startswith('wrn'): model = models.__dict__[args.arch]( num_classes=num_classes, depth=args.depth, widen_factor=args.widen_factor, dropRate=args.drop, ) elif args.arch.endswith('resnet'): model = models.__dict__[args.arch]( num_classes=num_classes, depth=args.depth, ) #model = torch.nn.DataParallel(model).cuda() model = model.to(device) cudnn.benchmark = True print(' Total params: %.2fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0)) criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # Resume title = 'fashionmnist-' + args.arch if args.resume: # Load checkpoint. print('==> Resuming from checkpoint..') assert os.path.isfile( args.resume), 'Error: no checkpoint directory found!' args.checkpoint = os.path.dirname(args.resume) checkpoint = torch.load(args.resume) best_acc = checkpoint['best_acc'] start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title, resume=True) else: logger = Logger(os.path.join(args.checkpoint, 'log.txt'), title=title) logger.set_names([ 'Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.', 'Valid Acc.' ]) if args.evaluate: print('\nEvaluation only') #test_loss, test_acc = test(testloader, model, criterion, start_epoch, use_cuda) #print(' Test Loss: %.8f, Test Acc: %.2f' % (test_loss, test_acc)) model.eval() id = 0 preds_list = [] with torch.no_grad(): for X, y in testloader: batch_pred = list( model(X.to(device)).argmax(dim=1).cpu().numpy()) for y_pred in batch_pred: preds_list.append((id, y_pred)) id += 1 print('生成提交结果文件') with open('submission.csv', 'w') as f: f.write('ID,Prediction\n') for id, pred in preds_list: f.write('{},{}\n'.format(id, pred)) return # Train and val for epoch in range(start_epoch, args.epochs): adjust_learning_rate(optimizer, epoch) print('\nEpoch: [%d | %d] LR: %f' % (epoch + 1, args.epochs, state['lr'])) train_loss, train_acc = train(trainloader, model, criterion, optimizer, epoch, use_cuda) test_loss, test_acc = test(testloader, model, criterion, epoch, use_cuda) # append logger file logger.append( [state['lr'], train_loss, test_loss, train_acc, test_acc]) # save model is_best = test_acc > best_acc best_acc = max(test_acc, best_acc) save_checkpoint( { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'acc': test_acc, 'best_acc': best_acc, 'optimizer': optimizer.state_dict(), }, is_best, checkpoint=args.checkpoint) logger.close() logger.plot() savefig(os.path.join(args.checkpoint, 'log.eps')) print('Best acc:') print(best_acc)
def main(): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print("Using {} device training.".format(device.type)) # 用来保存coco_info的文件 results_file = "results{}.txt".format( datetime.datetime.now().strftime("%Y%m%d-%H%M%S")) # 检查保存权重文件夹是否存在,不存在则创建 if not os.path.exists("save_weights"): os.makedirs("save_weights") data_transform = { "train": transforms.Compose( [transforms.ToTensor(), transforms.RandomHorizontalFlip(0.5)]), "val": transforms.Compose([transforms.ToTensor()]) } VOC_root = "./" # check voc root if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False: raise FileNotFoundError( "VOCdevkit dose not in path:'{}'.".format(VOC_root)) # load train data set # VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt train_data_set = VOC2012DataSet(VOC_root, data_transform["train"], "train.txt") # 注意这里的collate_fn是自定义的,因为读取的数据包括image和targets,不能直接使用默认的方法合成batch batch_size = 8 nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers print('Using %g dataloader workers' % nw) train_data_loader = torch.utils.data.DataLoader( train_data_set, batch_size=batch_size, shuffle=True, num_workers=nw, collate_fn=train_data_set.collate_fn) # load validation data set # VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt val_data_set = VOC2012DataSet(VOC_root, data_transform["val"], "val.txt") val_data_set_loader = torch.utils.data.DataLoader( val_data_set, batch_size=batch_size, shuffle=False, num_workers=nw, collate_fn=train_data_set.collate_fn) # create model num_classes equal background + 20 classes model = create_model(num_classes=21) # print(model) model.to(device) train_loss = [] learning_rate = [] val_map = [] # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # first frozen backbone and train 5 epochs # # 首先冻结前置特征提取网络权重(backbone),训练rpn以及最终预测网络部分 # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # for param in model.backbone.parameters(): param.requires_grad = False # define optimizer params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005) init_epochs = 5 for epoch in range(init_epochs): # train for one epoch, printing every 10 iterations mean_loss, lr = utils.train_one_epoch(model, optimizer, train_data_loader, device, epoch, print_freq=50) train_loss.append(mean_loss.item()) learning_rate.append(lr) # evaluate on the test dataset coco_info = utils.evaluate(model, val_data_set_loader, device=device) # write into txt with open(results_file, "a") as f: # 写入的数据包括coco指标还有loss和learning rate result_info = [ str(round(i, 4)) for i in coco_info + [mean_loss.item(), lr] ] txt = "epoch:{} {}".format(epoch, ' '.join(result_info)) f.write(txt + "\n") val_map.append(coco_info[1]) # pascal mAP torch.save(model.state_dict(), "./save_weights/pretrain.pth") # # # # # # # # # # # # # # # # # # # # # # # # # # # # # second unfrozen backbone and train all network # # 解冻前置特征提取网络权重(backbone),接着训练整个网络权重 # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 冻结backbone部分底层权重 for name, parameter in model.backbone.named_parameters(): split_name = name.split(".")[0] if split_name in ["0", "1", "2", "3"]: parameter.requires_grad = False else: parameter.requires_grad = True # define optimizer params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005) # learning rate scheduler lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.33) num_epochs = 20 for epoch in range(init_epochs, num_epochs + init_epochs, 1): # train for one epoch, printing every 50 iterations mean_loss, lr = utils.train_one_epoch(model, optimizer, train_data_loader, device, epoch, print_freq=50) train_loss.append(mean_loss.item()) learning_rate.append(lr) # update the learning rate lr_scheduler.step() # evaluate on the test dataset coco_info = utils.evaluate(model, val_data_set_loader, device=device) # write into txt with open(results_file, "a") as f: # 写入的数据包括coco指标还有loss和learning rate result_info = [ str(round(i, 4)) for i in coco_info + [mean_loss.item(), lr] ] txt = "epoch:{} {}".format(epoch, ' '.join(result_info)) f.write(txt + "\n") val_map.append(coco_info[1]) # pascal mAP # save weights # 仅保存最后5个epoch的权重 if epoch in range(num_epochs + init_epochs)[-5:]: save_files = { 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'lr_scheduler': lr_scheduler.state_dict(), 'epoch': epoch } torch.save(save_files, "./save_weights/mobile-model-{}.pth".format(epoch)) # plot loss and lr curve if len(train_loss) != 0 and len(learning_rate) != 0: from plot_curve import plot_loss_and_lr plot_loss_and_lr(train_loss, learning_rate) # plot mAP curve if len(val_map) != 0: from plot_curve import plot_map plot_map(val_map)
def main(): torch.backends.cudnn.benchmark = True args = parser.parse_args() transform = transforms.Compose([ transforms.ToTensor(), transforms.Resize((224, 224)), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) ### data loader => MJJ checkpoint = torch.load(args.resnet_checkpoint) state_dict = { str.replace(k, 'module.', ''): v for k, v in checkpoint['state_dict'].items() } fe = resnet.resnet50() fe.load_state_dict(state_dict, False) model = torch.nn.Sequential(fe, clsf()) checkpoint = (torch.load(args.clsf_checkpoint)) state_dict = { str.replace(k, 'module.', ''): v for k, v in checkpoint['model'].items() } model.load_state_dict(state_dict, False) model = torch.nn.DataParallel(model).cuda() model.eval() cap = cv2.VideoCapture(args.input_filename) fps = 1 n = 0 num_frames = 0 out_dict = [] check = 0 frame_counter = 0 images = [] with torch.no_grad(): while cap.isOpened(): ret, image = cap.read() frame_counter += 1 if frame_counter >= (cap.get(cv2.CAP_PROP_FPS) / fps): if not ret: check = 10 #End else: image = Image.fromarray(image[:, :, ::-1]) images.append(image) check += 1 #Append 10 times. if check == 10: check = 0 data = transform(images).unsqueeze(0).pin_memory().cuda( non_blocking=True) output = model( data) #data : B*T*C*H*W // B : 1 // T : 10 frames. y_friends = F.softmax(output, 1) top5_value_friends, top5_index_friends = y_friends.topk(5) top5_value_friends, top5_label_friends = top5_value_friends.tolist( ), [class_friends[i] for i in top5_index_friends[:, 0]] for i in range(len(images)): num_frames += 1 out_dict.append({ "type": "location", "class": top5_label_friends[i], "seconds": float(num_frames) * 1.0 / float(fps) }) images = [] if not ret: break frame_counter = 0 n += 1 if n % 100 == 0: print('Processed {}/{} frames'.format( n, int(cap.get(cv2.CAP_PROP_FRAME_COUNT)))) sys.stdout.flush() if n == 1e10: break with jsonlines.open(args.output_filename + '.jsonl', mode='w') as writer: writer.write_all(out_dict) cap.release() print('Done')
def load_data(test_data_path, preprocess_path, fold, batch_size, num_workers): test_data_path = '/data/xxx/LUNA/rowfile/subset' crop_size = 32 black_list = [] preprocess_path = '/data/xxx/LUNA/cls/crop_v3' pix_value, npix = 0, 0 for file_name in os.listdir(preprocess_path): if file_name.endswith('.npy'): if file_name[:-4] in black_list: continue data = np.load(os.path.join(preprocess_path, file_name)) pix_value += np.sum(data) npix += np.prod(data.shape) pix_mean = pix_value / float(npix) pix_value = 0 for file_name in os.listdir(preprocess_path): if file_name.endswith('.npy'): if file_name[:-4] in black_list: continue data = np.load(os.path.join(preprocess_path, file_name)) - pix_mean pix_value += np.sum(data * data) pix_std = np.sqrt(pix_value / float(npix)) print(pix_mean, pix_std) transform_train = transforms.Compose([ # transforms.RandomScale(range(28, 38)), transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.RandomYFlip(), transforms.RandomZFlip(), transforms.ZeroOut(4), transforms.ToTensor(), transforms.Normalize( (pix_mean), (pix_std)), # need to cal mean and std, revise norm func ]) transform_test = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((pix_mean), (pix_std)), ]) # load data list test_file_name_list = [] test_label_list = [] test_feat_list = [] data_frame = pd.read_csv('./data/annotationdetclsconvfnl_v3.csv', names=[ 'seriesuid', 'coordX', 'coordY', 'coordZ', 'diameter_mm', 'malignant' ]) all_list = data_frame['seriesuid'].tolist()[1:] label_list = data_frame['malignant'].tolist()[1:] crdx_list = data_frame['coordX'].tolist()[1:] crdy_list = data_frame['coordY'].tolist()[1:] crdz_list = data_frame['coordZ'].tolist()[1:] dim_list = data_frame['diameter_mm'].tolist()[1:] # test id test_id_list = [] for file_name in os.listdir(test_data_path + str(fold) + '/'): if file_name.endswith('.mhd'): test_id_list.append(file_name[:-4]) mxx = mxy = mxz = mxd = 0 for srsid, label, x, y, z, d in zip(all_list, label_list, crdx_list, crdy_list, crdz_list, dim_list): mxx = max(abs(float(x)), mxx) mxy = max(abs(float(y)), mxy) mxz = max(abs(float(z)), mxz) mxd = max(abs(float(d)), mxd) if srsid in black_list: continue # crop raw pixel as feature data = np.load(os.path.join(preprocess_path, srsid + '.npy')) bgx = int(data.shape[0] / 2 - crop_size / 2) bgy = int(data.shape[1] / 2 - crop_size / 2) bgz = int(data.shape[2] / 2 - crop_size / 2) data = np.array(data[bgx:bgx + crop_size, bgy:bgy + crop_size, bgz:bgz + crop_size]) y, x, z = np.ogrid[-crop_size / 2:crop_size / 2, -crop_size / 2:crop_size / 2, -crop_size / 2:crop_size / 2] mask = abs(y**3 + x**3 + z**3) <= abs(float(d))**3 feat = np.zeros((crop_size, crop_size, crop_size), dtype=float) feat[mask] = 1 if srsid.split('-')[0] in test_id_list: test_file_name_list.append(srsid + '.npy') test_label_list.append(int(label)) test_feat_list.append(feat) for idx in range(len(test_feat_list)): test_feat_list[idx][-1] /= mxd test_set = lunanod(preprocess_path, test_file_name_list, test_label_list, test_feat_list, train=False, download=True, transform=transform_test) test_loader = torch.utils.data.DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=num_workers) return test_loader
def main(args): init_distributed_mode(args) print(args) device = torch.device(args.device) # 用来保存coco_info的文件 results_file = "results{}.txt".format(datetime.datetime.now().strftime("%Y%m%d-%H%M%S")) # Data loading code print("Loading data") data_transform = { "train": transforms.Compose([transforms.ToTensor(), transforms.RandomHorizontalFlip(0.5)]), "val": transforms.Compose([transforms.ToTensor()]) } VOC_root = args.data_path # check voc root if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False: raise FileNotFoundError("VOCdevkit dose not in path:'{}'.".format(VOC_root)) # load train data set # VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt train_dataset = VOCDataSet(VOC_root, "2012", data_transform["train"], "train.txt") # load validation data set # VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt val_dataset = VOCDataSet(VOC_root, "2012", data_transform["val"], "val.txt") print("Creating data loaders") if args.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) test_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset) else: train_sampler = torch.utils.data.RandomSampler(train_dataset) test_sampler = torch.utils.data.SequentialSampler(val_dataset) if args.aspect_ratio_group_factor >= 0: # 统计所有图像比例在bins区间中的位置索引 group_ids = create_aspect_ratio_groups(train_dataset, k=args.aspect_ratio_group_factor) train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids, args.batch_size) else: train_batch_sampler = torch.utils.data.BatchSampler( train_sampler, args.batch_size, drop_last=True) data_loader = torch.utils.data.DataLoader( train_dataset, batch_sampler=train_batch_sampler, num_workers=args.workers, collate_fn=train_dataset.collate_fn) data_loader_test = torch.utils.data.DataLoader( val_dataset, batch_size=1, sampler=test_sampler, num_workers=args.workers, collate_fn=train_dataset.collate_fn) print("Creating model") # create model # 注意:不包含背景 model = create_model(num_classes=args.num_classes) model.to(device) model_without_ddp = model if args.distributed: model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu]) model_without_ddp = model.module params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD( params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.lr_step_size, gamma=args.lr_gamma) lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.lr_steps, gamma=args.lr_gamma) # 如果传入resume参数,即上次训练的权重地址,则接着上次的参数训练 if args.resume: # If map_location is missing, torch.load will first load the module to CPU # and then copy each parameter to where it was saved, # which would result in all processes on the same machine using the same set of devices. checkpoint = torch.load(args.resume, map_location='cpu') # 读取之前保存的权重文件(包括优化器以及学习率策略) model_without_ddp.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) args.start_epoch = checkpoint['epoch'] + 1 if args.test_only: utils.evaluate(model, data_loader_test, device=device) return train_loss = [] learning_rate = [] val_map = [] print("Start training") start_time = time.time() for epoch in range(args.start_epoch, args.epochs): if args.distributed: train_sampler.set_epoch(epoch) mean_loss, lr = utils.train_one_epoch(model, optimizer, data_loader, device, epoch, args.print_freq, warmup=True) train_loss.append(mean_loss.item()) learning_rate.append(lr) # update learning rate lr_scheduler.step() # evaluate after every epoch coco_info = utils.evaluate(model, data_loader_test, device=device) val_map.append(coco_info[1]) # pascal mAP # 只在主进程上进行写操作 if args.rank in [-1, 0]: # write into txt with open(results_file, "a") as f: # 写入的数据包括coco指标还有loss和learning rate result_info = [str(round(i, 4)) for i in coco_info + [mean_loss.item()]] + [str(round(lr, 6))] txt = "epoch:{} {}".format(epoch, ' '.join(result_info)) f.write(txt + "\n") if args.output_dir: # 只在主节点上执行保存权重操作 save_on_master({ 'model': model_without_ddp.state_dict(), 'optimizer': optimizer.state_dict(), 'lr_scheduler': lr_scheduler.state_dict(), 'args': args, 'epoch': epoch}, os.path.join(args.output_dir, 'model_{}.pth'.format(epoch))) total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str)) if args.rank in [-1, 0]: # plot loss and lr curve if len(train_loss) != 0 and len(learning_rate) != 0: from plot_curve import plot_loss_and_lr plot_loss_and_lr(train_loss, learning_rate) # plot mAP curve if len(val_map) != 0: from plot_curve import plot_map plot_map(val_map)
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_dataset(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 = False trainloader = DataLoader( VideoDataset(dataset.train, seq_len=args.seq_len, sample='random', transform=transform_train), sampler=RandomIdentitySampler(dataset.train, num_instances=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='dense', 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='dense', 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: {:.5f}M".format( sum(p.numel() for p in model.parameters()) / 1000000.0)) criterion_xent = CrossEntropyLabelSmooth( num_classes=dataset.num_train_pids, use_gpu=use_gpu) criterion_htri = TripletLoss(margin=args.margin) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) if args.stepsize > 0: scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma) start_epoch = args.start_epoch start_time = time.time() print(start_time) for batch_idx, (imgs, pids, _) in enumerate(trainloader): print(batch_idx) print('x') if use_gpu: imgs, pids = imgs.cuda(), pids.cuda() imgs, pids = Variable(imgs), Variable(pids) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
def main(args): print(args) # mp.spawn(main_worker, args=(args,), nprocs=args.world_size, join=True) utils.init_distributed_mode(args) device = torch.device(args.device) # Data loading code print("Loading data") data_transform = { "train": transforms.Compose( [transforms.ToTensor(), transforms.RandomHorizontalFlip(0.5)]), "val": transforms.Compose([transforms.ToTensor()]) } VOC_root = args.data_path # load train data set train_data_set = VOC2012DataSet(VOC_root, data_transform["train"], True) # load validation data set val_data_set = VOC2012DataSet(VOC_root, data_transform["val"], False) print("Creating data loaders") if args.distributed: train_sampler = torch.utils.data.distributed.DistributedSampler( train_data_set) test_sampler = torch.utils.data.distributed.DistributedSampler( val_data_set) else: train_sampler = torch.utils.data.RandomSampler(train_data_set) test_sampler = torch.utils.data.SequentialSampler(val_data_set) if args.aspect_ratio_group_factor >= 0: # 统计所有图像比例在bins区间中的位置索引 group_ids = create_aspect_ratio_groups( train_data_set, k=args.aspect_ratio_group_factor) train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids, args.batch_size) else: train_batch_sampler = torch.utils.data.BatchSampler(train_sampler, args.batch_size, drop_last=True) data_loader = torch.utils.data.DataLoader( train_data_set, batch_sampler=train_batch_sampler, num_workers=args.workers, collate_fn=utils.collate_fn) data_loader_test = torch.utils.data.DataLoader(val_data_set, batch_size=1, sampler=test_sampler, num_workers=args.workers, collate_fn=utils.collate_fn) print("Creating model") model = create_model(num_classes=21) model.to(device) model_without_ddp = model if args.distributed: model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[args.gpu]) model_without_ddp = model.module params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) # lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.lr_step_size, gamma=args.lr_gamma) lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones=args.lr_steps, gamma=args.lr_gamma) # 如果传入resume参数,即上次训练的权重地址,则接着上次的参数训练 if args.resume: # If map_location is missing, torch.load will first load the module to CPU # and then copy each parameter to where it was saved, # which would result in all processes on the same machine using the same set of devices. checkpoint = torch.load( args.resume, map_location='cpu') # 读取之前保存的权重文件(包括优化器以及学习率策略) model_without_ddp.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) args.start_epoch = checkpoint['epoch'] + 1 if args.test_only: utils.evaluate(model, data_loader_test, device=device) return print("Start training") start_time = time.time() for epoch in range(args.start_epoch, args.epochs): if args.distributed: train_sampler.set_epoch(epoch) utils.train_one_epoch(model, optimizer, data_loader, device, epoch, args.print_freq) lr_scheduler.step() if args.output_dir: # 只在主节点上执行保存权重操作 utils.save_on_master( { 'model': model_without_ddp.state_dict(), 'optimizer': optimizer.state_dict(), 'lr_scheduler': lr_scheduler.state_dict(), 'args': args, 'epoch': epoch }, os.path.join(args.output_dir, 'model_{}.pth'.format(epoch))) # evaluate after every epoch utils.evaluate(model, data_loader_test, device=device) total_time = time.time() - start_time total_time_str = str(datetime.timedelta(seconds=int(total_time))) print('Training time {}'.format(total_time_str))
if torch.cuda.is_available(): device = torch.device("cuda") else: device = torch.cuda("cpu") sys.stdout = Logger(osp.join(PATH, 'log_train.txt')) print("Dataset is being initialized") dataset = dataset_manager.init_img_dataset(root='data', name=dataset_name, split_id=0) tfms_train = tfms.Compose([ tfms.Random2DTranslation(256, 128), tfms.RandomHorizontalFlip(), tfms.ToTensor(), tfms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) tfms_test = tfms.Compose([ tfms.Resize((256, 128)), tfms.ToTensor(), tfms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) pin_memory = True trainloader = DataLoader( ImageDataset(dataset.train, transform=tfms_train), sampler=RandomIdentitySampler(dataset.train, num_instances=4), batch_size=32,
def get_transform(is_train): transforms = [] transforms.append(T.ToTensor()) return T.Compose(transforms)
def __init__(self): self.transforms = T.ToTensor()
import torch.nn as nn import torch.nn.functional as F import os from torch.autograd import Variable import transforms as transforms from skimage import io from skimage.transform import resize from models import * cut_size = 44 transform_test = transforms.Compose([ transforms.TenCrop(cut_size), transforms.Lambda(lambda crops: torch.stack( [transforms.ToTensor()(crop) for crop in crops])), ]) def rgb2gray(rgb): return np.dot(rgb[..., :3], [0.299, 0.587, 0.114]) raw_img = io.imread('images/1.jpg') gray = rgb2gray(raw_img) gray = resize(gray, (48, 48), mode='symmetric').astype(np.uint8) img = gray[:, :, np.newaxis] img = np.concatenate((img, img, img), axis=2) img = Image.fromarray(img)
learning_rate_decay_rate = 0.9 # 0.9 cut_size = 44 total_epoch = 250 path = os.path.join(opt.dataset + '_' + opt.model) dir_path = '/home/zjian30/test/Facial-Expression-Recognition.Pytorch/' path = dir_path + path print(path) # Data print('==> Preparing data..') transform_train = transforms.Compose([ transforms.RandomCrop(44), transforms.RandomHorizontalFlip(), transforms.ToTensor(), ]) transform_test = transforms.Compose([ transforms.TenCrop(cut_size), transforms.Lambda(lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops])), ]) trainset = FER2013(split = 'Training', transform=transform_train) trainloader = torch.utils.data.DataLoader(trainset, batch_size=opt.bs, shuffle=True, num_workers=1) PublicTestset = FER2013(split = 'PublicTest', transform=transform_test) PublicTestloader = torch.utils.data.DataLoader(PublicTestset, batch_size=opt.bs, shuffle=False, num_workers=1) PrivateTestset = FER2013(split = 'PrivateTest', transform=transform_test) PrivateTestloader = torch.utils.data.DataLoader(PrivateTestset, batch_size=opt.bs, shuffle=False, num_workers=1) # Model
def attr_main(): stale_step = args.stalestep runId = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S') args.save_dir = os.path.join(args.save_dir, runId) if not os.path.exists(args.save_dir): os.mkdir(args.save_dir) print(args.save_dir) 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('./log_train_' + runId + '.txt') else: sys.stdout = Logger('./log_test_' + runId + '.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_dataset(name=args.dataset, min_seq_len=args.seq_len, attr=True) args.attr_lens = [[], [9, 10, 2]] #dataset.attr_lens args.columns = ["upcolor", "downcolor", "gender"] #dataset.columns print("Initializing model: {}".format(args.arch)) # if args.arch == "resnet50ta_attr" or args.arch == "resnet50ta_attr_newarch": if args.arch == 'attr_resnet503d': model = models.init_model(name=args.arch, attr_lens=args.attr_lens, model_type=args.model_type, num_classes=dataset.num_train_pids, sample_width=args.width, sample_height=args.height, sample_duration=args.seq_len) torch.backends.cudnn.benchmark = False else: model = models.init_model(name=args.arch, attr_lens=args.attr_lens, model_type=args.model_type) print("Model size: {:.5f}M".format( sum(p.numel() for p in model.parameters()) / 1000000.0)) if args.dataset == "duke": 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]) ]) elif args.dataset == "mars": 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 # if args.predict: # attr_predict(model, transform_test, use_gpu) # return # # For SQL mmir, just first_img_path/ random_img_path sampling works # # For other mmir, call first_img/ random_img sampling # mmirImgQueryLoader = VideoDataset(dataset.mmir_img_query, seq_len=1, # dataset.mmir_img_query[0:2] # sample='random_img_path', # sample='random', #changing to random to get the imgs for running attr. extraction model, does the img_path match with 'random_img_path' sampling -- yap, the samples returned with both sampling matches in postgres and here, so next just get those images and run the attribute ex. models on them # transform=transform_test, attr=True, # attr_loss=args.attr_loss, attr_lens=args.attr_lens) # # The above function basically just tells how we sample dfrom the img list of this specific tracklet (has pid, camid, attrs) # # sampling more from gallery as test and gallery has same person, train has different persons # # during retrieval same person with same features should have higher rank than different person with same features # # 2/3 for image, rest is for video # mmirImgQueryLoader = VideoDataset( # dataset.mmir_img_gallery, # seq_len=1, # dataset.mmir_img_query[0:2] # sample='random_img_path', # sample='random', # transform=transform_test, attr=True, # attr_loss=args.attr_loss, attr_lens=args.attr_lens) # print(len(dataset.mmir_img_gallery)) # print(len(dataset.mmir_video_gallery)) # # import psycopg2 # # # set-up a postgres connection # # conn = psycopg2.connect(database='ng', user='******',password='******', # # host='146.148.89.5', port=5432) # # dbcur = conn.cursor() # # print("connection successful") # # for batch_idx, (pids, camids, attrs, img_path) in enumerate(tqdm(mmirImgQueryLoader)): # # #print(pids) # # images = list() # # images.append(img_path) # # # just save the img in img_path # # sql = dbcur.mogrify(""" # # INSERT INTO mmir_ground ( # # mgid, # # ctype, # # pid, # # camid, # # ubc, # # lbc, # # gender, # # imagepaths # # ) VALUES (DEFAULT, %s,%s,%s,%s,%s,%s,%s) ON CONFLICT DO NOTHING;""", ("image", # # str(pids), # # str(camids), # # str(attrs[0]), # # str(attrs[1]), # # str(attrs[2]), images) # # ) # # # print(sql) # # dbcur.execute(sql) # # conn.commit() # # print(dataset.mmir_video_query[0]) # if this works, for sdml or other mmir training will just take some part of train for img, some for video # # For SQL mmir, just first_img_path sampling works, as not even using the image_array, using just the attributes # # For other mmir, call random sampling # # sampling the whole video has no impact in mmir, as we cannot take avg. of the separated lists attributes "accuracy" # # we need frames of whose attributes would be one single one, and then we search in img and text that attributes # mmirVideoQueryLoader = VideoDataset(dataset.mmir_video_query, seq_len=10, # 100 * 10 = 1000 frames as image # sample='random_video', transform=transform_test, attr=True, # attr_loss=args.attr_loss, attr_lens=args.attr_lens) # mmirVideoQueryLoader = VideoDataset(dataset.mmir_video_gallery, seq_len=6, # 100 * 10 = 1000 frames as image # sample='random_video', transform=transform_test, attr=True, # attr_loss=args.attr_loss, attr_lens=args.attr_lens) # # for batch_idx, (pids, camids, attrs, img_paths) in enumerate(tqdm(mmirVideoQueryLoader)): # # #print(pids) # # sql = dbcur.mogrify(""" # # INSERT INTO mmir_ground ( # # mgid, # # ctype, # # pid, # # camid, # # ubc, # # lbc, # # gender, # # imagepaths # # ) VALUES (DEFAULT, %s,%s,%s,%s,%s,%s,%s) ON CONFLICT DO NOTHING;""", ("video", # # str(pids), # # str(camids), # # str(attrs[0]), # # str(attrs[1]), # # str(attrs[2]), img_paths) # # ) # # dbcur.execute(sql) # # conn.commit() # # conn.close() # # if i want to add more mmir image, video samples, # # 1. in line 116, pass num_mmir_query_imgs=1000, num_mmir_query_videos=100 # # 2. just uncomment the SQL query code, the first 1000 & 100 would be same, as seed is set, the later ones will be added in postgres # # hahahah, way surpassed above steps, created new samples from query_id, appended them; then created samples from gallery_id, added them by calling sampler again # trainloader = DataLoader( # VideoDataset(dataset.train, seq_len=args.seq_len, sample='random', transform=transform_train, attr=True, # attr_loss=args.attr_loss, attr_lens=args.attr_lens), # sampler=RandomIdentitySampler(dataset.train, num_instances=args.num_instances), # batch_size=args.train_batch, num_workers=args.workers, # pin_memory=pin_memory, drop_last=True, # ) # validloader = VideoDataset(dataset.valid, seq_len=args.seq_len, # sample='dense', transform=transform_test, attr=True, # attr_loss=args.attr_loss, attr_lens=args.attr_lens) # #queryloader = VideoDataset(dataset.query + dataset.gallery, seq_len=args.seq_len, sample='single' transform=transform_test, attr=True, # queryloader = VideoDataset(dataset.query + dataset.gallery, seq_len=args.seq_len, #args.seq_len, sample='dense' # sample='dense', transform=transform_test, attr=True, # attr_loss=args.attr_loss, attr_lens=args.attr_lens) # # queryLoaderMIT = VideoDataset(dataset.query + dataset.gallery, seq_len=1, # # sample='random_img_path', transform=transform_test, attr=True, # # attr_loss=args.attr_loss, attr_lens=args.attr_lens) start_epoch = args.start_epoch if use_gpu: model = nn.DataParallel(model).cuda() start_time = time.time() if args.arch == 'resnet503d': torch.backends.cudnn.benchmark = False # print("Run attribute pre-training") if args.attr_loss == "cropy": criterion = nn.CrossEntropyLoss() elif args.attr_loss == "mse": criterion = nn.MSELoss() if args.colorsampling: attr_test_MIT(model, criterion, queryLoaderMIT, use_gpu) return if args.evaluate: print("Evaluate only") # model_root = "/data/chenzy/models/mars/2019-02-26_21-02-13" model_root = args.eval_model_root model_paths = [] for m in os.listdir(model_root): if m.endswith("pth"): model_paths.append(m) model_paths = sorted(model_paths, key=lambda a: float(a.split("_")[1]), reverse=True) #print(model_paths) #input() # model_paths = ['rank1_2.8755379380596713_checkpoint_ep500.pth'] for m in model_paths[0:1]: model_path = os.path.join(model_root, m) print(model_path) old_weights = torch.load(model_path) new_weights = model.module.state_dict() for k in new_weights: if k in old_weights: new_weights[k] = old_weights[k] model.module.load_state_dict(new_weights) if args.predict: attr_predict(model, transform_test, use_gpu) return avr_acc = attr_test(model, criterion, queryloader, use_gpu) # break # test(model, queryloader, galleryloader, args.pool, use_gpu) return if use_gpu: optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.module.parameters()), lr=args.lr, weight_decay=args.weight_decay) else: optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay) # avr_acc = attr_test(model, criterion, queryloader, use_gpu) if args.stepsize > 0: scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma) best_avr = 0 no_rise = 0 for epoch in range(start_epoch, args.max_epoch): print("==> Epoch {}/{}".format(epoch + 1, args.max_epoch)) attr_train(model, criterion, optimizer, trainloader, use_gpu) if args.stepsize > 0: scheduler.step() if args.eval_step > 0 and ((epoch + 1) % (args.eval_step) == 0 or (epoch + 1) == args.max_epoch): # avr_acc = attr_test(model, criterion, queryloader, use_gpu) avr_acc = attr_test(model, criterion, validloader, use_gpu) print("avr", avr_acc) if avr_acc > best_avr: no_rise = 0 print("==> Test") best_avr = avr_acc if use_gpu: state_dict = model.module.state_dict() else: state_dict = model.state_dict() torch.save( state_dict, osp.join( args.save_dir, "avr_" + str(avr_acc) + '_checkpoint_ep' + str(epoch + 1) + '.pth')) else: no_rise += 1 print("no_rise:", no_rise) if no_rise > stale_step: break elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))
def __init__(self, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)): self.transforms = T.Compose([ T.ToTensor(), T.Normalize(mean=mean, std=std), ])
def main(): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print("Using {} device training.".format(device.type)) # 检查保存权重文件夹是否存在,不存在则创建 if not os.path.exists("save_weights"): os.makedirs("save_weights") data_transform = { "train": transforms.Compose( [transforms.ToTensor(), transforms.RandomHorizontalFlip(0.5)]), "val": transforms.Compose([transforms.ToTensor()]) } VOC_root = "../../data_set/" # check voc root if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False: raise FileNotFoundError( "VOCdevkit dose not in path:'{}'.".format(VOC_root)) # load train data set train_data_set = VOC2012DataSet(VOC_root, data_transform["train"], True) # 注意这里的collate_fn是自定义的,因为读取的数据包括image和targets,不能直接使用默认的方法合成batch batch_size = 4 nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers print('Using %g dataloader workers' % nw) train_data_loader = torch.utils.data.DataLoader( train_data_set, batch_size=batch_size, shuffle=True, num_workers=nw, collate_fn=train_data_set.collate_fn) # load validation data set val_data_set = VOC2012DataSet(VOC_root, data_transform["val"], False) val_data_set_loader = torch.utils.data.DataLoader( val_data_set, batch_size=batch_size, shuffle=False, num_workers=nw, collate_fn=train_data_set.collate_fn) # create model num_classes equal background + 20 classes model = create_model(num_classes=21) # print(model) model.to(device) train_loss = [] learning_rate = [] val_mAP = [] # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # first frozen backbone and train 5 epochs # # 首先冻结前置特征提取网络权重(backbone),训练rpn以及最终预测网络部分 # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # for param in model.backbone.parameters(): param.requires_grad = False # define optimizer params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005) num_epochs = 5 for epoch in range(num_epochs): # train for one epoch, printing every 10 iterations utils.train_one_epoch(model, optimizer, train_data_loader, device, epoch, print_freq=50, train_loss=train_loss, train_lr=learning_rate) # evaluate on the test dataset utils.evaluate(model, val_data_set_loader, device=device, mAP_list=val_mAP) torch.save(model.state_dict(), "./save_weights/pretrain.pth") # # # # # # # # # # # # # # # # # # # # # # # # # # # # # second unfrozen backbone and train all network # # 解冻前置特征提取网络权重(backbone),接着训练整个网络权重 # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # 冻结backbone部分底层权重 for name, parameter in model.backbone.named_parameters(): split_name = name.split(".")[0] if split_name in ["0", "1", "2", "3"]: parameter.requires_grad = False else: parameter.requires_grad = True # define optimizer params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005) # learning rate scheduler lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.33) num_epochs = 20 for epoch in range(num_epochs): # train for one epoch, printing every 50 iterations utils.train_one_epoch(model, optimizer, train_data_loader, device, epoch, print_freq=50, train_loss=train_loss, train_lr=learning_rate) # update the learning rate lr_scheduler.step() # evaluate on the test dataset utils.evaluate(model, val_data_set_loader, device=device, mAP_list=val_mAP) # save weights if epoch > 10: save_files = { 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'lr_scheduler': lr_scheduler.state_dict(), 'epoch': epoch } torch.save(save_files, "./save_weights/mobile-model-{}.pth".format(epoch)) # plot loss and lr curve if len(train_loss) != 0 and len(learning_rate) != 0: from plot_curve import plot_loss_and_lr plot_loss_and_lr(train_loss, learning_rate) # plot mAP curve if len(val_mAP) != 0: from plot_curve import plot_map plot_map(val_mAP)
def main(parser_data): device = torch.device( parser_data.device if torch.cuda.is_available() else "cpu") print("Using {} device training.".format(device.type)) # 用来保存coco_info的文件 results_file = "results{}.txt".format( datetime.datetime.now().strftime("%Y%m%d-%H%M%S")) data_transform = { "train": transforms.Compose( [transforms.ToTensor(), transforms.RandomHorizontalFlip(0.5)]), "val": transforms.Compose([transforms.ToTensor()]) } VOC_root = parser_data.data_path # check voc root if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False: raise FileNotFoundError( "VOCdevkit dose not in path:'{}'.".format(VOC_root)) # load train data set # VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt train_data_set = VOC2012DataSet(VOC_root, data_transform["train"], "train.txt") # 注意这里的collate_fn是自定义的,因为读取的数据包括image和targets,不能直接使用默认的方法合成batch batch_size = parser_data.batch_size nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers print('Using %g dataloader workers' % nw) train_data_loader = torch.utils.data.DataLoader( train_data_set, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=nw, collate_fn=train_data_set.collate_fn) # load validation data set # VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt val_data_set = VOC2012DataSet(VOC_root, data_transform["val"], "val.txt") val_data_set_loader = torch.utils.data.DataLoader( val_data_set, batch_size=batch_size, shuffle=False, pin_memory=True, num_workers=nw, collate_fn=train_data_set.collate_fn) # create model num_classes equal background + 20 classes model = create_model(num_classes=parser_data.num_classes + 1, device=device) # print(model) model.to(device) # define optimizer params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005) # learning rate scheduler lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.33) # 如果指定了上次训练保存的权重文件地址,则接着上次结果接着训练 if parser_data.resume != "": checkpoint = torch.load(parser_data.resume, map_location=device) model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) parser_data.start_epoch = checkpoint['epoch'] + 1 print("the training process from epoch{}...".format( parser_data.start_epoch)) train_loss = [] learning_rate = [] val_map = [] for epoch in range(parser_data.start_epoch, parser_data.epochs): # train for one epoch, printing every 10 iterations mean_loss, lr = utils.train_one_epoch(model, optimizer, train_data_loader, device=device, epoch=epoch, print_freq=50, warmup=True) train_loss.append(mean_loss.item()) learning_rate.append(lr) # update the learning rate lr_scheduler.step() # evaluate on the test dataset coco_info = utils.evaluate(model, val_data_set_loader, device=device) # write into txt with open(results_file, "a") as f: # 写入的数据包括coco指标还有loss和learning rate result_info = [ str(round(i, 4)) for i in coco_info + [mean_loss.item()] ] + [str(round(lr, 6))] txt = "epoch:{} {}".format(epoch, ' '.join(result_info)) f.write(txt + "\n") val_map.append(coco_info[1]) # pascal mAP # save weights save_files = { 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'lr_scheduler': lr_scheduler.state_dict(), 'epoch': epoch } torch.save(save_files, "./save_weights/resNetFpn-model-{}.pth".format(epoch)) # plot loss and lr curve if len(train_loss) != 0 and len(learning_rate) != 0: from plot_curve import plot_loss_and_lr plot_loss_and_lr(train_loss, learning_rate) # plot mAP curve if len(val_map) != 0: from plot_curve import plot_map plot_map(val_map)
def from_image(cls, img, sr_factor): pil = transforms.ToPILImage()(transforms.ToTensor()(img)) return ZSSRDataset(pil, sr_factor)
def main(): SIZE_IMG = 1024; # 224 global args args = parser.parse_args() # create model if args.arch.startswith('resnet50') : model = model_defs.resnet50_oneway(num_classes=2); elif args.arch.startswith('vgg16_bn') : model = model_defs.vgg16_bn_oneway(num_classes=2, pretrained=True); elif args.arch.startswith('segnet') : model = model_defs.segnet(num_classes=2, pretrained=True); model = nn.DataParallel(model.model); # dirty trick #model = nn.DataParallel(model); # dirty trick # open log file if args.train == 1 : log_dir = 'logs'; log_name = args.arch + '_new.csv'; if not os.path.isdir(log_dir) : os.mkdir(log_dir); log_handle = get_file_handle(os.path.join(log_dir, log_name), 'wb+'); log_handle.write('Epoch, LearningRate, Momentum, WeightDecay,' + \ 'Loss, Precision, Recall, Accuracy(IoU), FgWeight, BgWeight\n'); log_handle.close(); # check model directory model_dir = 'models'; if not os.path.isdir(model_dir) : os.mkdir(model_dir); # resume learning based on cmdline arguments if ((args.start_epoch > 1) and (args.train==1)) : load_epoch = args.start_epoch - 1; elif (args.train==0) : load_epoch = args.load_epoch; else : load_epoch = 0; if load_epoch > 0 : print("=> loading checkpoint for epoch = '{}'" .format(load_epoch)); checkpoint_name = args.arch + '_ep_' + str(load_epoch) + '.pth.tar'; checkpoint = torch.load(os.path.join(model_dir, checkpoint_name)); model.load_state_dict(checkpoint['state_dict']); # if args.arch.startswith('resnet50') : # model = add_dropout2d(model); # if args.arch.startswith('vgg16_bn') : # model = add_dropout2d_vgg(model); #model = add_conv_layers_vgg(model); model.cuda(); # transfer to cuda print(model); mean = load_pickle('./mean'); std = load_pickle('./std'); if args.train == 1 : train_data_dir, train_gt_dir = args.data, args.gt; train_loader = torch.utils.data.DataLoader( datasets.ImageFolder(train_data_dir, train_gt_dir, transform_joint=transforms.Compose_Joint([ transforms.RandomCrop(SIZE_IMG), transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(), ]), transform=transforms.Compose([ transforms.ColorJitter(0.3, 0.3, 0.3, 0), transforms.ToTensor(), transforms.Normalize(mean=mean, std=std),]), target_transform=transforms.Compose([ transforms.ToTensorTarget(),]), do_copy=True, ), batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True); weights = torch.from_numpy(np.array([1.,1.2])).float(); criterion = nn.CrossEntropyLoss(weights).cuda(); if args.optim == 'adam' : optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay); elif args.optim == 'sgd' : optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay); for epoch in range(args.start_epoch, args.end_epoch+1): # train for one epoch stats_epoch = train(train_loader, model, criterion, optimizer, epoch); model_name = args.arch + '_ep_' + str(epoch) + '.pth.tar'; # get current parameters of optimizer for param_group in optimizer.param_groups : cur_lr = param_group['lr']; cur_wd = param_group['weight_decay']; if param_group.has_key('momentum') : cur_momentum = param_group['momentum']; else : cur_momentum = 'n/a'; break; # constant parameters throughout the network if epoch % args.save_interval == 0 : state = { 'epoch': epoch, 'arch': args.arch, 'state_dict': model.state_dict(), 'learning_rate': cur_lr, 'moemntum': cur_momentum, 'weight_decay': cur_wd, 'fg_weight': weights[1], 'bg_weight': weights[0], }; torch.save(state, os.path.join(model_dir, model_name)); # write logs using logHandle log_handle = get_file_handle(os.path.join(log_dir, log_name), 'ab'); log_handle.write(str(epoch) + ',' + str(cur_lr) + ',' + str(cur_momentum) + ',' + str(cur_wd) + ',' + str(stats_epoch['loss']) + ',' + str(stats_epoch['prec']) + ',' + str(stats_epoch['recall']) + ',' + str(stats_epoch['acc']) + ',' + str(weights[1]) + ',' + str(weights[0]) + '\n'); log_handle.close(); # adjust_learning_rate(optimizer, epoch, 10); # adjust learning rate elif args.train == 0 : # test testdir = args.data; outdir = args.out; # stride = args.test_stride; # test_batch_size = args.test_batch_size; test_transformer = transforms.Compose([ # transforms.RandomHorizontalFlip(), # transforms.RandomVerticalFlip(), transforms.ToTensor(), transforms.Normalize(mean=mean, std=std),]); # test(testdir, outdir, test_transformer, model, load_epoch, stride, SIZE_IMG); # test_batch_form(testdir, outdir, test_transformer, model, load_epoch, # stride, SIZE_IMG, test_batch_size); test_full_res(testdir, outdir, test_transformer, model, load_epoch, SIZE_IMG);
def get_transform(train): transforms = [] transforms.append(T.ToTensor()) if train: transforms.append(T.RandomHorizontalFlip(0.5)) return T.Compose(transforms)
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_dataset(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_train_p = T.Compose([ T.Random2DTranslation(256, 128), 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]), ]) transform_test_p = T.Compose([ T.Resize((256, 128)), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) train_file = 'data/cuhk_train.pkl' test_file = 'data/cuhk_test.pkl' gallery_file = 'data/cuhk_gallery.pkl' data_root = args.data_root dataset_train = CUHKGroup(train_file, data_root, True, transform_train, transform_train_p) dataset_test = CUHKGroup(test_file, data_root, False, transform_test, transform_test_p) dataset_query = CUHKGroup(test_file, data_root, False, transform_test, transform_test_p) dataset_gallery = CUHKGroup(gallery_file, data_root, False, transform_test, transform_test_p) pin_memory = True if use_gpu else False if args.xent_only: trainloader = DataLoader( dataset_train, batch_size=args.train_batch, shuffle=True, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) else: trainloader = DataLoader( dataset_train, sampler=RandomIdentitySampler(dataset_train, num_instances=args.num_instances), batch_size=args.train_batch, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) queryloader = DataLoader( dataset_test, batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) querygalleryloader = DataLoader( dataset_query, batch_size=args.gallery_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) galleryloader = DataLoader( dataset_gallery, batch_size=args.gallery_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=True, ) print("Initializing model: {}".format(args.arch)) if args.xent_only: # model = models.init_model(name=args.arch, num_classes=dataset_train.num_train_gids, loss={'xent'}) model = models.init_model(name=args.arch, num_classes=dataset_train.num_train_gids, loss={'xent'}) else: # model = models.init_model(name=args.arch, num_classes=dataset_train.num_train_gids, loss={'xent', 'htri'}) model = models.init_model( name=args.arch, num_classes=dataset_train.num_train_gids, num_person_classes=dataset_train.num_train_pids, loss={'xent', 'htri'}) #criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset_train.num_train_gids, use_gpu=use_gpu) #criterion_xent_person = CrossEntropyLabelSmooth(num_classes=dataset_train.num_train_pids, use_gpu=use_gpu) if os.path.exists(args.pretrained_model): print("Loading checkpoint from '{}'".format(args.pretrained_model)) checkpoint = torch.load(args.pretrained_model) model_dict = model.state_dict() pretrain_dict = checkpoint['state_dict'] pretrain_dict = { k: v for k, v in pretrain_dict.items() if k in model_dict } model_dict.update(pretrain_dict) model.load_state_dict(model_dict) criterion_xent = nn.CrossEntropyLoss(ignore_index=-1) criterion_xent_person = nn.CrossEntropyLoss(ignore_index=-1) criterion_htri = TripletLoss(margin=args.margin) criterion_pair = ContrastiveLoss(margin=args.margin) criterion_htri_filter = TripletLossFilter(margin=args.margin) criterion_permutation = PermutationLoss() optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr, weight_decay=args.weight_decay) if args.stepsize > 0: if args.warmup: scheduler = WarmupMultiStepLR(optimizer, [200, 400, 600]) else: scheduler = lr_scheduler.StepLR(optimizer, step_size=args.stepsize, gamma=args.gamma) start_epoch = args.start_epoch if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") test_gcn_person_batch(model, queryloader, querygalleryloader, galleryloader, args.pool, use_gpu) #test_gcn_batch(model, queryloader, querygalleryloader, galleryloader, args.pool, use_gpu) #test_gcn(model, queryloader, galleryloader, args.pool, use_gpu) #test(model, queryloader, galleryloader, args.pool, use_gpu) return start_time = time.time() best_rank1 = -np.inf for epoch in range(start_epoch, args.max_epoch): #print("==> Epoch {}/{} lr:{}".format(epoch + 1, args.max_epoch, scheduler.get_lr()[0])) train_gcn(model, criterion_xent, criterion_xent_person, criterion_pair, criterion_htri_filter, criterion_htri, criterion_permutation, optimizer, trainloader, use_gpu) #train(model, criterion_xent, criterion_htri, optimizer, trainloader, use_gpu) if args.stepsize > 0: scheduler.step() if args.eval_step > 0 and (epoch + 1) % args.eval_step == 0 or ( epoch + 1) == args.max_epoch: print("==> Test") rank1 = test_gcn_person_batch(model, queryloader, querygalleryloader, galleryloader, args.pool, use_gpu) #rank1 = test_gcn(model, queryloader, galleryloader, args.pool, use_gpu=False) #rank1 = test(model, queryloader, galleryloader, args.pool, use_gpu) is_best = rank1 > best_rank1 if is_best: best_rank1 = rank1 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')) elapsed = round(time.time() - start_time) elapsed = str(datetime.timedelta(seconds=elapsed)) print("Finished. Total elapsed time (h:m:s): {}".format(elapsed))