def __init__(self, annFile: str, imageDir: str, targetHeight: int, targetWidth: int, numClass: int, train: bool = True): self.annotations = {} self.table = { '不良-機械傷害': 0, '不良-著色不佳': 1, '不良-炭疽病': 2, '不良-乳汁吸附': 3, '不良-黑斑病': 4 } self.imageDir = imageDir self.numClass = numClass with open(annFile, 'r', encoding='utf-8-sig') as f: for line in f.readlines(): arr = line.rstrip().split(',') ans = [] for idx in range(1, len(arr), 5): tlx, tly, w, h, c = arr[idx:idx + 5] if tlx: tlx, tly, w, h = list(map(float, (tlx, tly, w, h))) if c not in self.table: self.table[c] = len(self.table) cx = tlx + w / 2 cy = tly + h / 2 c = self.table[c] ans.append(list(map(int, (cx, cy, w, h, c)))) self.annotations[arr[0]] = ans self.names = list(self.annotations) with open('table.txt', 'w') as f: f.write(str(self.table)) print(self.table) if train: self.transforms = T.Compose([ T.RandomOrder([ T.RandomHorizontalFlip(), T.RandomVerticalFlip(), T.RandomSizeCrop(numClass) ]), T.Resize((targetHeight, targetWidth)), T.ColorJitter(brightness=.2, contrast=0, saturation=0, hue=0), T.Normalize() ]) else: self.transforms = T.Compose( [T.Resize((targetHeight, targetWidth)), T.Normalize()])
def __getitem__(self, index): if self.is_train: img, target = self.train_img[index], self.train_label[index] if len(img.shape) == 2: img = np.stack([img] * 3, 2) img = Image.fromarray(img, mode="RGB") img = transforms.Resize((256, 256), Image.BILINEAR)(img) img = transforms.RandomCrop(INPUT_SIZE)(img) img = transforms.RandomHorizontalFlip()(img) img = transforms.ToTensor()(img) img = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(img) else: img, target = self.test_img[index], self.test_label[index] if len(img.shape) == 2: img = np.stack([img] * 3, 2) img = Image.fromarray(img, mode="RGB") img = transforms.Resize((256, 256), Image.BILINEAR)(img) img = transforms.CenterCrop(INPUT_SIZE)(img) img = transforms.ToTensor()(img) img = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(img) return img, target
def __init__(self, cms, margin_size, margin_time): self.IMAGE_SCALE = 20 self.IMAGE_SIZE = self.IMAGE_SCALE * self.IMAGE_SCALE self.train_transform = transforms.Compose([ transforms.Resize((self.IMAGE_SCALE, self.IMAGE_SCALE)), transforms.ToTensor() ]) self.test_transform = transforms.Compose([ transforms.Resize((self.IMAGE_SCALE, self.IMAGE_SCALE)), transforms.ToTensor() ]) self.CMS = cms self.MARGIN_SIZE = margin_size self.MARGIN_TIME = margin_time
def __init__(self, root: str, annotation: str, targetHeight: int, targetWidth: int, numClass: int): self.root = root self.coco = COCO(annotation) self.ids = list(self.coco.imgs.keys()) self.targetHeight = targetHeight self.targetWidth = targetWidth self.numClass = numClass self.transforms = T.Compose([ T.RandomOrder( [T.RandomHorizontalFlip(), T.RandomSizeCrop(numClass)]), T.Resize((targetHeight, targetWidth)), T.ColorJitter(brightness=.2, contrast=.1, saturation=.1, hue=0), T.Normalize() ]) self.newIndex = {} classes = [] for i, (k, v) in enumerate(self.coco.cats.items()): self.newIndex[k] = i classes.append(v['name']) with open('classes.txt', 'w') as f: f.write(str(classes))
def __init__(self, root: str, targetHeight: int, targetWidth: int, numClass: int, train: bool = True): """ :param root: should contain .jpg files and corresponding .txt files :param targetHeight: desired height for model input :param targetWidth: desired width for model input :param numClass: number of classes in the given dataset """ self.cache = {} imagePaths = glob(os.path.join(root, '*.jpg')) for path in imagePaths: name = path.split('/')[-1].split('.jpg')[0] self.cache[path] = os.path.join(root, f'{name}.txt') self.paths = list(self.cache.keys()) self.targetHeight = targetHeight self.targetWidth = targetWidth self.numClass = numClass if train: self.transforms = T.Compose([ T.RandomOrder([ T.RandomHorizontalFlip(), T.RandomVerticalFlip(), T.RandomSizeCrop(numClass) ]), T.Resize((targetHeight, targetWidth)), T.ColorJitter(brightness=.2, contrast=.1, saturation=.1, hue=0), T.Normalize() ]) else: self.transforms = T.Compose( [T.Resize((targetHeight, targetWidth)), T.Normalize()])
def __init__(self, root=None, dataloader=default_loader): self.transform1 = transforms.Compose([ transforms.RandomRotation(30), transforms.Resize([256, 256]), transforms.RandomCrop(INPUT_SIZE), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness=(0.9, 1.1), contrast=(0.9, 1.1), saturation=(0.9, 1.1)), transforms.ToTensor(), transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), transforms.RandomErasing(probability=0.5, sh=0.05) ]) # 增强方法2: 关注更小的区域 self.transform2 = transforms.Compose([ transforms.RandomRotation(30), transforms.Resize([336, 336]), transforms.RandomCrop(INPUT_SIZE), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness=(0.9, 1.1), contrast=(0.9, 1.1), saturation=(0.9, 1.1)), transforms.ToTensor(), transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), transforms.RandomErasing(probability=0.5, sh=0.05) ]) self.dataloader = dataloader self.root = root with open(os.path.join(self.root, TRAIN_DATASET), 'r') as fid: self.imglist = fid.readlines() self.labels = [] for line in self.imglist: image_path, label = line.strip().split() self.labels.append(int(label)) self.labels = np.array(self.labels) self.labels = torch.LongTensor(self.labels)
def __init__(self, root=None, dataloader=default_loader): self.transform = transforms.Compose([ transforms.Resize([256, 256]), transforms.RandomCrop(INPUT_SIZE), transforms.ToTensor(), transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)) ]) self.dataloader = dataloader self.root = root self.imgs = [] self.labels = [] with open(os.path.join(self.root, EVAL_DATASET), 'r') as fid: for line in fid.readlines(): img_path, label = line.strip().split() img = self.dataloader(img_path) label = int(label) self.imgs.append(img) self.labels.append(label)
def get_transform(train=False, yolo=False, aug=None): assert aug == 'dirty_camera_lens' or aug == 'gan' or aug is None, "Aug parameter not valid" transforms = [] if yolo: transforms.append(custom_T.PadToSquare()) transforms.append(custom_T.Resize(img_size=None)) if train: transforms.append(custom_T.RandomHorizontalFlip()) transforms.append(custom_T.RandomCrop()) if aug == 'dirty_camera_lens': print("Augmentation: Dirty Camera Lens") transforms.append(custom_T.DirtyCameraLens()) transforms.append(custom_T.ToTensor()) # transforms.append(custom_T.FasterRCNNResizer()) return custom_T.Compose(transforms)
def deploy(path): assert os.path.exists(path), f'{path} not found : (' dataset = 'YOUR_DATASET_NAME' img_size = 256 test_transform = transforms.Compose([ transforms.Resize((img_size, img_size)), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) testA = ImageFolder(os.path.join('dataset', dataset, 'testA'), test_transform) with fluid.dygraph.guard(): testA_loader = DataLoader(testA, batch_size=1, shuffle=False) real_A, _ = next(iter(testA_loader)) in_np = real_A.numpy() # load model place = fluid.CPUPlace() exe = fluid.Executor(place) program, feed_vars, fetch_vars = fluid.io.load_inference_model(path, exe) # inference fetch, = exe.run(program, feed={feed_vars[0]: in_np}, fetch_list=fetch_vars) def img_postprocess(img): assert isinstance(img, np.ndarray), type(img) img = img * 0.5 + 0.5 img = img.squeeze(0).transpose((1, 2, 0)) # BGR to RGB img = img[:, :, ::-1] return img in_img = img_postprocess(in_np) out_img = img_postprocess(fetch) plt.subplot(121) plt.title('real A') plt.imshow(in_img) plt.subplot(122) plt.title('A to B') plt.imshow(out_img) plt.show()
with open('data/kannada_semi_1pct.pkl', 'rb') as f: kannada_semi = pickle.load(f) kannada_x_train_labeled = kannada_semi["x_train_labeled"] kannada_y_train_labeled = kannada_semi["y_train_labeled"] kannada_x_train_unlabeled = kannada_semi["x_train_unlabeled"] kannada_y_train_unlabeled = kannada_semi["y_train_unlabeled"] kannada_x_train = np.concatenate((kannada_x_train_labeled, kannada_x_train_unlabeled), axis=0) kannada_y_train = np.concatenate((kannada_y_train_labeled, kannada_y_train_unlabeled), axis=0) kannada_x_val = kannada_semi["x_val"] kannada_y_val = kannada_semi["y_val"] kannada_x_test = kannada_semi['x_test'] kannada_y_test = kannada_semi['y_test'] train_transform = transforms.Compose([ transforms.Normalize(mean=0.5, std=0.5), transforms.Resize(size=(32, 32)) ]) mnist_train_dataset = CustomTensorDataset(torch.from_numpy(mnist_x_train).float(), torch.from_numpy(mnist_y_train).long(), transform=train_transform) kannada_train_dataset = CustomTensorDataset(torch.from_numpy(kannada_x_train).float(), torch.from_numpy(kannada_y_train).long(), transform=train_transform) kannada_val_dataset = CustomTensorDataset(torch.from_numpy(kannada_x_val).float(), torch.from_numpy(kannada_y_val).long(), transform=train_transform) train_loader_1 = DataLoader(mnist_train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) train_loader_2 = DataLoader(kannada_train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) val_loader = DataLoader(kannada_val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers) num_batch = min(len(train_loader_1), len(train_loader_2)) dim_z = 512 # Modules
def main(args=None): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices # torch.cuda.set_device(0) 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, args.log_train)) else: sys.stdout = Logger(osp.join(args.save_dir, args.log_test)) 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.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), T.RandomErasing(), ]) transform_train2 = T.Compose([ T.Resize((args.height, args.width)), T.Random2DTranslation(args.height, args.width), ]) 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, data_name=args.dataset, seq_len=args.seq_len, sample='random', transform=transform_train, transform2=transform_train2, type="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, data_name=args.dataset, seq_len=args.seq_len, sample='dense', transform=transform_test, type="test"), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( VideoDataset(dataset.gallery, data_name=args.dataset, seq_len=args.seq_len, sample='dense', transform=transform_test, type="test"), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing models: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, final_dim=args.feat_dim) print("Model size: {:.5f}M".format( sum(p.numel() for p in model.parameters()) / 1000000.0)) crossEntropyLoss = CrossEntropyLabelSmooth( num_classes=dataset.num_train_pids, use_gpu=use_gpu) tripletLoss = TripletLoss(margin=args.margin) regularLoss = RegularLoss(use_gpu=use_gpu) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) scheduler = WarmupMultiStepLR(optimizer, args.stepsize, args.gamma, args.warmup_factor, args.warmup_items, args.warmup_method) start_epoch = args.start_epoch if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") atest(model, queryloader, galleryloader, use_gpu) return start_time = time.time() best_rank1 = -np.inf for epoch in range(start_epoch, args.max_epoch): print("==> Epoch {}/{}".format(epoch + 1, args.max_epoch)) train(model, crossEntropyLoss, tripletLoss, regularLoss, optimizer, trainloader, use_gpu) # if args.stepsize > 0: scheduler.step() if (epoch + 1) >= 200 and (epoch + 1) % args.eval_step == 0: print("==> Test") rank1 = atest(model, queryloader, galleryloader, 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, }, is_best, osp.join( args.save_dir, args.model_name + 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))
def get_data(data_dir, source, target, source_train_path, target_train_path, source_extension, target_extension, height, width, batch_size, re=0, workers=8): dataset = DA(data_dir, source, target, source_train_path, target_train_path, source_extension, target_extension) normalizer = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) source_num_classes = dataset.num_source_train_ids train_transformer = T.Compose([ T.RandomSizedRectCrop(height, width), T.RandomHorizontalFlip(), T.ToTensor(), normalizer, T.RandomErasing(EPSILON=re), ]) test_transformer = T.Compose([ T.Resize((height, width), interpolation=3), T.ToTensor(), normalizer, ]) source_train_loader = DataLoader(Preprocessor( dataset.source_train, root=osp.join(dataset.source_images_dir, dataset.source_train_path), transform=train_transformer), batch_size=batch_size, num_workers=0, shuffle=True, pin_memory=False, drop_last=True) target_train_loader = DataLoader(Preprocessor( dataset.target_train, root=osp.join(dataset.target_images_dir, dataset.target_train_path), transform=train_transformer), batch_size=batch_size, num_workers=0, shuffle=True, pin_memory=False, drop_last=True) # source_train_loader = DataLoader( # UnsupervisedCamStylePreprocessor(dataset.source_train, root=osp.join(dataset.source_images_dir, dataset.source_train_path), # camstyle_root=osp.join(dataset.source_images_dir, dataset.source_train_path), # transform=train_transformer), # batch_size=batch_size, num_workers=0, # shuffle=True, pin_memory=False, drop_last=True) # target_train_loader = DataLoader( # UnsupervisedCamStylePreprocessor(dataset.target_train, # root=osp.join(dataset.target_images_dir, dataset.target_train_path), # camstyle_root=osp.join(dataset.target_images_dir, # dataset.target_train_camstyle_path), # num_cam=dataset.target_num_cam, transform=train_transformer), # batch_size=batch_size, num_workers=workers, # shuffle=True, pin_memory=True, drop_last=True) query_loader = DataLoader(Preprocessor(dataset.query, root=osp.join( dataset.target_images_dir, dataset.query_path), transform=test_transformer), batch_size=batch_size, num_workers=workers, shuffle=False, pin_memory=True) gallery_loader = DataLoader(Preprocessor(dataset.gallery, root=osp.join( dataset.target_images_dir, dataset.gallery_path), transform=test_transformer), batch_size=batch_size, num_workers=workers, shuffle=False, pin_memory=True) return dataset, source_num_classes, source_train_loader, target_train_loader, query_loader, gallery_loader
def main(args=None): torch.manual_seed(args.seed) os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_devices # torch.cuda.set_device(0) 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, args.log_train)) else: sys.stdout = Logger(osp.join(args.save_dir, args.log_test)) 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_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 queryloader = DataLoader( VideoDataset(dataset.query, data_name=args.dataset, seq_len=args.seq_len, sample='dense', transform=transform_test, type="test"), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) galleryloader = DataLoader( VideoDataset(dataset.gallery, data_name=args.dataset, seq_len=args.seq_len, sample='dense', transform=transform_test, type="test"), batch_size=args.test_batch, shuffle=False, num_workers=args.workers, pin_memory=pin_memory, drop_last=False, ) print("Initializing models: {}".format(args.arch)) model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, final_dim=args.feat_dim) print("Model size: {:.5f}M".format( sum(p.numel() for p in model.parameters()) / 1000000.0)) checkpoint = torch.load(args.model_path) model.load_state_dict(checkpoint['state_dict']) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") atest(model, queryloader, galleryloader, use_gpu) return
def build_model(self): """ DataLoader """ pad = int(30 * self.img_size // 256) train_transform = transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.Resize((self.img_size + pad, self.img_size + pad)), transforms.RandomCrop(self.img_size), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) test_transform = transforms.Compose([ transforms.Resize((self.img_size, self.img_size)), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) self.trainA = ImageFolder( os.path.join('dataset', self.dataset, 'trainA'), train_transform) self.trainB = ImageFolder( os.path.join('dataset', self.dataset, 'trainB'), train_transform) self.testA = ImageFolder( os.path.join('dataset', self.dataset, 'testA'), test_transform) self.testB = ImageFolder( os.path.join('dataset', self.dataset, 'testB'), test_transform) self.trainA_loader = DataLoader(self.trainA, batch_size=self.batch_size, shuffle=True) self.trainB_loader = DataLoader(self.trainB, batch_size=self.batch_size, shuffle=True) self.testA_loader = DataLoader(self.testA, batch_size=1, shuffle=False) self.testB_loader = DataLoader(self.testB, batch_size=1, shuffle=False) """ Define Generator, Discriminator """ self.genA2B = ResnetGenerator(input_nc=3, output_nc=3, ngf=self.ch, n_blocks=self.n_res, img_size=self.img_size, light=self.light) self.genB2A = ResnetGenerator(input_nc=3, output_nc=3, ngf=self.ch, n_blocks=self.n_res, img_size=self.img_size, light=self.light) self.disGA = Discriminator(input_nc=3, ndf=self.ch, n_layers=7) self.disGB = Discriminator(input_nc=3, ndf=self.ch, n_layers=7) self.disLA = Discriminator(input_nc=3, ndf=self.ch, n_layers=5) self.disLB = Discriminator(input_nc=3, ndf=self.ch, n_layers=5) """ Define Loss """ self.L1_loss = loss.L1Loss() self.MSE_loss = loss.MSELoss() self.BCE_loss = loss.BCEWithLogitsLoss() """ Trainer """ def get_params(block): out = [] for name, param in block.named_parameters(): if 'instancenorm' in name or 'weight_u' in name or 'weight_v' in name: continue out.append(param) return out genA2B_parameters = get_params(self.genA2B) genB2A_parameters = get_params(self.genB2A) disGA_parameters = get_params(self.disGA) disGB_parameters = get_params(self.disGB) disLA_parameters = get_params(self.disLA) disLB_parameters = get_params(self.disLB) G_parameters = genA2B_parameters + genB2A_parameters D_parameters = disGA_parameters + disGB_parameters + disLA_parameters + disLB_parameters self.G_optim = fluid.optimizer.Adam( parameter_list=G_parameters, learning_rate=self.lr, beta1=0.5, beta2=0.999, regularization=fluid.regularizer.L2Decay(self.weight_decay)) self.D_optim = fluid.optimizer.Adam( parameter_list=D_parameters, learning_rate=self.lr, beta1=0.5, beta2=0.999, regularization=fluid.regularizer.L2Decay(self.weight_decay)) """ Define Rho clipper to constraint the value of rho in AdaILN and ILN"""
save_path = os.path.join(args.save_model,model_path) torch.save(tosave_model.module.state_dict(), save_path) with open(os.path.join(args.save_model,'checkpoint.txt'),'w') as fin: fin.write(model_path + ' ' + str(epoch) + '\n') #dataset prepare #--------------------------------- print('Loading dataset...') cache_size = 256 if args.image_size == 448: cache_size = 256 * 2 if args.image_size == 352: cache_size = 402 transform_train = transforms.Compose([ transforms.Resize((cache_size,cache_size)), #transforms.Resize((args.image_size,args.image_size)), #transforms.RandomRotation(10), transforms.RandomCrop(args.image_size), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485,0.456,0.406], [0.229,0.224,0.225]), ]) transform_test = transforms.Compose([ transforms.Resize((cache_size,cache_size)), transforms.CenterCrop(args.image_size), transforms.ToTensor(), transforms.Normalize([0.485,0.456,0.406], [0.229,0.224,0.225]), ])
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, ) 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), 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( ImageDataset(dataset.test_3000_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.test_3000_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.train_vehicle_nums, loss={'xent', 'htri'}) print("Model size: {:.3f} M".format(sum(p.numel() for p in model.parameters()) / 1000000.0)) modelid = 0 if args.upload: modelid = upload_data.updatemodel(args.arch, args.lr, args.stepsize, args.gamma, args.loss, args.dataset, args.height, args.width, args.seq_len, args.train_batch, args.other) # criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset.train_vehicle_nums, use_gpu=use_gpu) 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) start_epoch = args.start_epoch if args.load_weights: # load pretrained weights but ignore layers that don't match in size print("Loading pretrained weights from '{}'".format(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) if args.resume: checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) start_epoch = checkpoint['epoch'] rank1 = checkpoint['rank1'] print("Loaded checkpoint from '{}'".format(args.resume)) print("- start_epoch: {}\n- rank1: {}".format(start_epoch, rank1)) if use_gpu: model = nn.DataParallel(model).cuda() if args.evaluate: print("Evaluate only") test(model, queryloader, galleryloader, use_gpu, return_distmat=True) 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(modelid, 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") cmc, mAP = test(model, queryloader, galleryloader, use_gpu) rank1 = cmc[0] 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')) if args.upload: upload_data.updatetest(modelid, epoch + 1, mAP.item(), cmc[0].item(), cmc[4].item(), cmc[9].item(), cmc[19].item()) upload_data.updaterank(modelid, best_rank1.item()) 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 generate_dataloader(model, eval_type, args): cache_size = 256 if args.image_size == 448: cache_size = 256 * 2 if args.image_size == 352: cache_size = 402 transform_train = transforms.Compose([ transforms.Resize((cache_size, cache_size)), transforms.RandomCrop(args.image_size), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) transform_test = transforms.Compose([ transforms.Resize((cache_size, cache_size)), transforms.CenterCrop(args.image_size), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) if eval_type == 'train': data = ImageDataset(max_person=args.max_person + 1, image_dir=args.images_root, images_list=args.train_file_pre + '_images.txt', bboxes_list=args.train_file_pre + '_bbox.json', image_size=args.image_size, input_transform=transform_train) elif eval_type == 'valid': data = ImageDataset(max_person=args.max_person + 1, image_dir=args.images_root, \ images_list=args.valid_file_pre + '_images.txt', \ bboxes_list=args.valid_file_pre + '_bbox.json', \ image_size=args.image_size, \ input_transform=transform_test) else: data = ImageDataset(max_person=args.max_person + 1, image_dir=args.images_root, \ images_list=args.test_file_pre + '_images.txt', \ bboxes_list=args.test_file_pre + '_bbox.json', \ image_size=args.image_size, \ input_transform=transform_test) loader = torch.utils.data.DataLoader(data, batch_size=args.test_batch_size, shuffle=False, num_workers=0, worker_init_fn=np.random.seed( args.manualSeed)) model.eval() batch_time = AverageMeter.AverageMeter() end_time = time.time() union_filename = './para_data/' + args.dataset + '_' + eval_type + '_union' + '.npy' feat_filename = './para_data/' + args.dataset + '_' + eval_type + '_ssl.npy' ssl_model = load_ssl_model() ssl_model.eval() if args.cuda: ssl_model.cuda() ssl_model = torch.nn.DataParallel(ssl_model) if os.path.exists(feat_filename) and os.path.exists( union_filename) and not args.regenerate_roifeat: all_feat = np.load(feat_filename) logger.info('loadding RoI feature npy from {} successfully'.format( feat_filename)) all_union_feat = np.load(union_filename, mmap_mode='r') # [B, N, N, 2048] logger.info( 'loading union npy from {} successfully'.format(union_filename)) else: all_feat, all_union_feat = [], [] for batch_idx, (img, image_bboxes) in enumerate(loader): if args.cuda: img, image_bboxes = img.cuda(), image_bboxes.cuda() img, image_bboxes = Variable(img), Variable(image_bboxes) node_num = image_bboxes.shape[1] union_boxes = get_union_box(image_bboxes) if args.cuda: union_boxes = union_boxes.cuda() image_bboxes = torch.cat((image_bboxes, union_boxes), axis=1) del union_boxes image_bboxes = Variable(image_bboxes) # [batcn, node_num, 2048] rois_feature_all = model(img, image_bboxes) feature_num = rois_feature_all.shape[2] rois_feature = rois_feature_all[:, :node_num, :] union_feature = rois_feature_all[:, node_num:, :].reshape( -1, node_num, node_num, feature_num) if args.load_ssl_model: img_feat = ssl_model( img, image_bboxes) # [batch_size, max_person+1, feat_dim] rois_feature = torch.cat((rois_feature, img_feat[:, -1:, :]), dim=1) all_feat.append(rois_feature.cpu().data.numpy()) all_union_feat.append(union_feature.cpu().data.numpy()) batch_time.update(time.time() - end_time) end_time = time.time() if batch_idx % args.print_freq == 0: logger.info('%s Epoch: [%d/%d] ' 'Time %.3f (%.3f)\t' % (eval_type, batch_idx, len(loader), batch_time.val, batch_time.avg)) all_feat = np.concatenate(all_feat) all_union_feat = np.concatenate(all_union_feat) np.save(feat_filename, all_feat) np.save(union_filename, all_union_feat) class_weight, class_count = [], [] if eval_type == 'train': dataset = SRDataset(all_feat, all_union_feat, max_person=args.max_person+1, image_dir=args.images_root, \ images_list=args.train_file_pre + '_images.txt', relations_list=args.train_file_pre + '_relation.json', image_size=args.image_size) class_weight, class_count = dataset.class_weight() batch_size = args.batch_size is_shuffle = True elif eval_type == 'valid': dataset = SRDataset(all_feat, all_union_feat, max_person=args.max_person+1,image_dir=args.images_root, \ images_list=args.valid_file_pre + '_images.txt', relations_list=args.valid_file_pre + '_relation.json', image_size=args.image_size) batch_size = args.test_batch_size is_shuffle = False else: dataset = SRDataset(all_feat, all_union_feat, max_person=args.max_person+1,image_dir=args.images_root, \ images_list=args.test_file_pre + '_images.txt', \ relations_list=args.test_file_pre + '_relation.json', image_size=args.image_size ) batch_size = args.test_batch_size is_shuffle = False dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=is_shuffle, num_workers=args.num_workers, worker_init_fn=np.random.seed( args.manualSeed)) return dataloader, class_weight, class_count
#torch.cuda.manual_seed(args.seed) kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {} ### PARAMETERS ### # LSTM & Q Learning IMAGE_SCALE = 28 IMAGE_SIZE = IMAGE_SCALE*IMAGE_SCALE HIDDEN_LAYERS = 1 HIDDEN_NODES = 200 OUTPUT_CLASSES = args.class_vector_size ################## train_transform = transforms.Compose([ transforms.Resize((IMAGE_SCALE, IMAGE_SCALE)), transforms.ToTensor() ]) test_transform = transforms.Compose([ transforms.Resize((IMAGE_SCALE, IMAGE_SCALE)), transforms.ToTensor() ]) print("Loading trainingsets...") omniglot_loader = loader.OmniglotLoader('data/omniglot', classify=False, partition=0.8, classes=True) train_loader = torch.utils.data.DataLoader( OMNIGLOT('data/omniglot', train=True, transform=train_transform, download=True, omniglot_loader=omniglot_loader, batch_size=args.episode_size), batch_size=args.mini_batch_size, shuffle=True, **kwargs) print("Loading testset...") test_loader = torch.utils.data.DataLoader( OMNIGLOT('data/omniglot', train=False, transform=test_transform, omniglot_loader=omniglot_loader, batch_size=args.episode_size),
def main(): # Initialize arguments args = init_arguments() # Initialize random init_seed(args.seed) # Initialize the model model, parameters, losses, accuracy_dict, accuracies, start_epoch = init_model( args) IMAGE_SCALE = 20 transform = transforms.Compose( [transforms.Resize((IMAGE_SCALE, IMAGE_SCALE)), transforms.ToTensor()]) omniglot_loader = loader.OmniglotLoader('data/omniglot', classify=False, partition=0.8, classes=True) train_loader = torch.utils.data.DataLoader( OMNIGLOT('data/omniglot', train=True, transform=transform, download=True, omniglot_loader=omniglot_loader, batch_size=parameters.episode_length), batch_size=parameters.batch_size, shuffle=True) test_loader = torch.utils.data.DataLoader(OMNIGLOT( 'data/omniglot', train=False, transform=transform, download=False, omniglot_loader=omniglot_loader, batch_size=parameters.episode_length), batch_size=parameters.batch_size, shuffle=False) criterion = nn.CrossEntropyLoss(reduce=False) optimizer = optim.Adam(model.parameters()) init_logging() accuracy_dict, losses, accuracies = train_model(model, parameters, train_loader, criterion, optimizer, args, accuracy_dict, losses, accuracies, start_epoch) loss_plot.plot([accuracies], ["Training Accuracy Percentage"], "training_stats", args.name + "/", "Percentage") loss_plot.plot([losses], ["Training Loss"], "training_loss", args.name + "/", "Average Loss") predictions, labels = evaluate(model, test_loader, criterion, parameters, args, 10) matrix_plot.plot_matrix(predictions, labels, args.name + "/", title="matrix_plot_test") scatterplot.plot(accuracy_dict, args.name + "/", parameters.batch_size, title="Prediction Accuracy after Testing")