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 load_dataloader(opt): if opt.imglist_file is not None: dataset = ImageList(opt.imglist_file, transform= \ transforms.Compose([DEFAULT_TRANSFORMS, Resize(opt.img_size)])) else: dataset = ImageFolder(opt.image_folder, transform= \ transforms.Compose([DEFAULT_TRANSFORMS, Resize(opt.img_size)])) dataloader = DataLoader(dataset, batch_size=opt.batch_size, shuffle=False, num_workers=opt.n_cpu, ) return dataloader
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 main(): val_transform_det = trans.Compose([ trans.Scale(256, 256), ]) val_data = dates.Dataset('/dataset', '/dataset', '/dataset/test.txt', 'val', transform=True, transform_med=val_transform_det) val_loader = Data.DataLoader(val_data, batch_size=1, shuffle=False, num_workers=4, pin_memory=True) import model.siameseNet.dares as models model = models.SiameseNet(norm_flag='l2') checkpoint = torch.load('the path to best model', map_location='cpu') model.load_state_dict(checkpoint['state_dict']) print('load success') model = model.cuda() save_change_map_dir = 'the path to changemap' save_roc_dir = 'the path to roc' time_start = time.time() current_metric = validate(model, val_loader, save_change_map_dir, save_roc_dir) elapsed = round(time.time() - time_start) elapsed = str(datetime.timedelta(seconds=elapsed)) print('Elapsed {}'.format(elapsed))
def load_data(args): normalize = t.Normalize(mean=[0.445, 0.287, 0.190], std=[0.31, 0.225, 0.168]) im_transform = t.Compose([t.ToTensor(), normalize]) # Use the following code fo co_transformations e.g. random rotation or random flip etc. # co_transformer = cot.Compose([cot.RandomRotate(45)]) dsetTrain = GIANA(args.imgdir, args.gtdir, input_size=(args.input_width, args.input_height), train=True, transform=im_transform, co_transform=None, target_transform=t.ToLabel()) train_data_loader = data.DataLoader(dsetTrain, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) dsetVal = GIANA(args.imgdir, args.gtdir, train=False, transform=im_transform, co_transform=None, target_transform=t.ToLabel()) val_data_loader = data.DataLoader(dsetVal, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers) return train_data_loader, val_data_loader
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, image_dir, crop_size=64, blend_mode="linear", horizontal_flip=False): """ :param image_dir (str): path of the images :param crop_size (int or tuple): crop size, default is 64 :param blend_mode (str): pick on of two, `screen` or `linear`, represents image composition type :param horizontal_flip (bool): Whether use horizontal flipping or not """ # super(DataLoader, self).__init__() # 1. initialize file path or a list of file names. assert blend_mode in ["screen", "linear"] self.blend_mode = blend_mode self.data_path = image_dir self.all_filenames = os.listdir(self.data_path) self.label_filenames = list( filter(lambda filename: filename.startswith("norain"), self.all_filenames)) self.num_files = len(self.label_filenames) print("[DataLoader] preprocess {} files on dir `{}`".format( self.num_files, self.data_path)) self.transform = transforms.Compose([ transforms.FiveCrop( crop_size, horizontal_flip), # tuple (tl, tr, bl, br, center) lambda crops: np.stack( [transforms.ToArray()(crop) for crop in crops]) ])
def loading_data(root, mode, batch_size=1): mean_std = ([0.5, 0.5, 0.5], [0.25, 0.25, 0.25]) log_para = 1 if mode == 'train': main_transform = own_transforms.Compose( [own_transforms.RandomHorizontallyFlip()]) else: main_transform = None img_transform = standard_transforms.Compose([ standard_transforms.ToTensor(), standard_transforms.Normalize(*mean_std) ]) gt_transform = standard_transforms.Compose( [own_transforms.LabelNormalize(log_para)]) restore_transform = standard_transforms.Compose([ own_transforms.DeNormalize(*mean_std), standard_transforms.ToPILImage() ]) dataset = HeadCountDataset(root, mode, main_transform, img_transform, gt_transform) if mode == 'train': dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True) else: dataloader = DataLoader(dataset, batch_size=1, shuffle=False) return dataloader, restore_transform
def get_transform(): transforms = [] transforms.append(custom_T.ToTensor()) transforms.append(custom_T.FasterRCNNResizer()) return custom_T.Compose(transforms)
def detect_image(model, image, img_size=416, conf_thres=0.5, nms_thres=0.5): """Inferences one image with model. :param model: Model for inference :type model: models.Darknet :param image: Image to inference :type image: nd.array :param img_size: Size of each image dimension for yolo, defaults to 416 :type img_size: int, optional :param conf_thres: Object confidence threshold, defaults to 0.5 :type conf_thres: float, optional :param nms_thres: IOU threshold for non-maximum suppression, defaults to 0.5 :type nms_thres: float, optional :return: Detections on image :rtype: nd.array """ model.eval() # Set model to evaluation mode # Configure input input_img = transforms.Compose([DEFAULT_TRANSFORMS, Resize(img_size)])((image, np.zeros( (1, 5))))[0].unsqueeze(0) if torch.cuda.is_available(): input_img = input_img.to("cuda") # Get detections with torch.no_grad(): detections = model(input_img) detections = non_max_suppression(detections, conf_thres, nms_thres) detections = rescale_boxes(detections[0], img_size, image.shape[:2]) return to_cpu(detections).numpy()
def get_coco(root, image_set, transforms, mode='instances'): anno_file_template = "{}_{}2017.json" PATHS = { "train": ("train2017", os.path.join("annotations", anno_file_template.format(mode, "train"))), "val": ("val2017", os.path.join("annotations", anno_file_template.format(mode, "val"))), # "train": ("val2017", os.path.join("annotations", anno_file_template.format(mode, "val"))) } t = [ConvertCocoPolysToMask()] if transforms is not None: t.append(transforms) transforms = T.Compose(t) img_folder, ann_file = PATHS[image_set] img_folder = os.path.join(root, img_folder) ann_file = os.path.join(root, ann_file) dataset = CocoDetection(img_folder, ann_file, transforms=transforms) if image_set == "train": dataset = _coco_remove_images_without_annotations(dataset) # dataset = torch.utils.data.Subset(dataset, [i for i in range(500)]) return dataset
def main(): ###### load datasets ######## train_transform_det = trans.Compose([ trans.Scale(cfg.TRANSFROM_SCALES), ]) val_transform_det = trans.Compose([ trans.Scale(cfg.TRANSFROM_SCALES), ]) train_data = data.Dataset(cfg.TRAIN_DATA_PATH, cfg.TRAIN_LABEL_PATH, cfg.TRAIN_TXT_PATH, 'train', transform=True, transform_med=train_transform_det) train_loader = Data.DataLoader(train_data, batch_size=8, shuffle=True, num_workers=4, pin_memory=True) val_data = data.Dataset(cfg.VAL_DATA_PATH, cfg.VAL_LABEL_PATH, cfg.VAL_TXT_PATH, 'val', transform=True, transform_med=val_transform_det) val_loader = Data.DataLoader(val_data, batch_size=cfg.BATCH_SIZE, shuffle=False, num_workers=4, pin_memory=True) ''' test_data = dates.Dataset(cfg.TEST_DATA_PATH,cfg.VAL_LABEL_PATH, cfg.VAL_TXT_PATH,'val',transform=True, transform_med = val_transform_det) test_loader = Data.DataLoader(test_data, batch_size= cfg.BATCH_SIZE, shuffle= False, num_workers= 4, pin_memory= True) ''' ###### build models ######## import model.DASNET as models G = models.generator D = models.discriminator train_model(D, G, train_loader, val_loader)
def get_transform(train): transforms = [] # converts the image, a PIL image, into a PyTorch Tensor transforms.append(T.ToTensor()) if train: # during training, randomly flip the training images # and ground-truth for data augmentation transforms.append(T.RandomHorizontalFlip(0.5)) return T.Compose(transforms)
def get_voc(root, image_set, transforms): t = [ConvertVOCtoCOCO()] if transforms is not None: t.append(transforms) transforms = T.Compose(t) dataset = VOCDetection(img_folder=root, year='2007', image_set=image_set, transforms=transforms) return dataset
def train(args): joint_transform = transforms.Compose([ transforms.RandomScale(), transforms.Mirror(), transforms.RandomCrop() ]) trainset = datasets[args.dataset](mode=args.mode, root=args.dataset_root) net = models[args.g]
def get_transform(train=False): transforms = [] if train: transforms.append(custom_T.RandomHorizontalFlip()) transforms.append(custom_T.RandomCrop()) transforms.append(custom_T.ToTensor()) transforms.append(custom_T.FasterRCNNResizer()) return custom_T.Compose(transforms)
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 prepare_transforms(tranform_params): transforms_list = [] for transform_info in tranform_params: transform_name = transform_info['name'] transform_params = transform_info['params'] if transform_params is not None: transform = transforms.__dict__[transform_name](**transform_params) else: transform = transforms.__dict__[transform_name]() transforms_list.append(transform) transform = transforms.Compose(transforms_list) return transform
def get_transforms(args): transforms_list = [ custom_transforms.RandomSizedCrop(args.image_size, args.resize_min, args.resize_max), custom_transforms.RandomHorizontalFlip(), custom_transforms.toTensor() ] if args.color_space == 'lab': transforms_list.insert(2, custom_transforms.toLAB()) elif args.color_space == 'rgb': transforms_list.insert(2, custom_transforms.toRGB('RGB')) transforms = custom_transforms.Compose(transforms_list) return transforms
def visualization(x_train, y_train, x_val, y_val, x_test, y_test, nb_class, ckpt, opt): # no augmentations used for linear evaluation transform_lineval = transforms.Compose([transforms.ToTensor()]) train_set_lineval = UCR2018(data=x_train, targets=y_train, transform=transform_lineval) val_set_lineval = UCR2018(data=x_val, targets=y_val, transform=transform_lineval) test_set_lineval = UCR2018(data=x_test, targets=y_test, transform=transform_lineval) train_loader_lineval = torch.utils.data.DataLoader(train_set_lineval, batch_size=128, shuffle=True) val_loader_lineval = torch.utils.data.DataLoader(val_set_lineval, batch_size=128, shuffle=False) test_loader_lineval = torch.utils.data.DataLoader(test_set_lineval, batch_size=128, shuffle=False) signal_length = x_train.shape[1] # loading the saved backbone backbone_lineval = SimConv4().cuda() # defining a raw backbone model checkpoint = torch.load(ckpt, map_location='cpu') backbone_lineval.load_state_dict(checkpoint) print('Visualization') for epoch in range(opt.epochs_test): backbone_lineval.eval() embeds = None labels = None for i, (data, target) in enumerate(test_loader_lineval): data = data.cuda() target = target.cuda().view(-1, 1) output = backbone_lineval(data).detach() if embeds is None: embeds = output.cpu().numpy() labels = target.cpu().numpy() else: embeds = np.vstack([embeds, output.cpu().numpy()]) labels = np.vstack([labels, target.cpu().numpy()]) return embeds, labels
def pretrain_InterSampleRel(x_train, y_train, opt): K = opt.K batch_size = opt.batch_size # 128 has been used in the paper tot_epochs = opt.epochs # 400 has been used in the paper feature_size = opt.feature_size ckpt_dir = opt.ckpt_dir prob = 0.2 # Transform Probability raw = transforms.Raw() cutout = transforms.Cutout(sigma=0.1, p=prob) jitter = transforms.Jitter(sigma=0.2, p=prob) scaling = transforms.Scaling(sigma=0.4, p=prob) magnitude_warp = transforms.MagnitudeWrap(sigma=0.3, knot=4, p=prob) time_warp = transforms.TimeWarp(sigma=0.2, knot=8, p=prob) window_slice = transforms.WindowSlice(reduce_ratio=0.8, p=prob) window_warp = transforms.WindowWarp(window_ratio=0.3, scales=(0.5, 2), p=prob) transforms_list = {'jitter': [jitter], 'cutout': [cutout], 'scaling': [scaling], 'magnitude_warp': [magnitude_warp], 'time_warp': [time_warp], 'window_slice': [window_slice], 'window_warp': [window_warp], 'G0': [jitter, magnitude_warp, window_slice], 'G1': [jitter, time_warp, window_slice], 'G2': [jitter, time_warp, window_slice, window_warp, cutout], 'none': [raw]} transforms_targets = list() for name in opt.aug_type: for item in transforms_list[name]: transforms_targets.append(item) train_transform = transforms.Compose(transforms_targets + [transforms.ToTensor()]) backbone = SimConv4().cuda() model = RelationalReasoning(backbone, feature_size).cuda() train_set = MultiUCR2018(data=x_train, targets=y_train, K=K, transform=train_transform) train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True) torch.save(model.backbone.state_dict(), '{}/backbone_init.tar'.format(ckpt_dir)) acc_max, epoch_max = model.train(tot_epochs=tot_epochs, train_loader=train_loader, opt=opt) torch.save(model.backbone.state_dict(), '{}/backbone_last.tar'.format(ckpt_dir)) return acc_max, epoch_max
def get_dataloader(): # TODO(xwd): Adaptive normalization by some large image. # E.g. In medical image processing, WSI image is very large and different to ordinary images. value_scale = 255 mean = [0.485, 0.456, 0.406] mean = [item * value_scale for item in mean] std = [0.229, 0.224, 0.225] std = [item * value_scale for item in std] train_transform = transform.Compose([ transform.RandomScale([cfg['scale_min'], cfg['scale_max']]), transform.RandomRotate([cfg['rotate_min'], cfg['rotate_max']], padding=mean, ignore_label=cfg['ignore_label']), transform.RandomGaussianBlur(), transform.RandomHorizontallyFlip(), transform.RandomCrop([cfg['train_h'], cfg['train_w']], crop_type='rand', padding=mean, ignore_label=cfg['ignore_label']), transform.ToTensor(), transform.Normalize(mean=mean, std=std) ]) train_data = cityscapes.Cityscapes(cfg['data_path'], split='train', transform=train_transform) # Use data sampler to make sure each GPU loads specific parts of dataset to avoid data reduntant. train_sampler = DistributedSampler(train_data) train_loader = DataLoader(train_data, batch_size=cfg['batch_size'] // cfg['world_size'], shuffle=(train_sampler is None), num_workers=4, pin_memory=True, sampler=train_sampler, drop_last=True) return train_loader, train_sampler
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 __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 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()
def main(): net = AFENet(classes=21, pretrained_model_path=None).cuda() net.load_state_dict(torch.load(os.path.join(args['model_save_path'], args['snapshot']))) value_scale = 255 mean = [0.485, 0.456, 0.406] mean = [item * value_scale for item in mean] std = [0.229, 0.224, 0.225] std = [item * value_scale for item in std] test_transform = transform.Compose([transform.ToTensor()]) test_data = voc2012.VOC2012(split='test', data_root=args['dataset_root'], data_list=args['test_list'], transform=test_transform) test_loader = DataLoader(test_data, batch_size=1, shuffle=False, num_workers=1) gray_folder = os.path.join(args['test_result_save_path']) # gray_folder = os.path.join(args['test_result_save_path'], 'gray') color_folder = os.path.join(args['test_result_save_path'], 'color') colors = np.loadtxt(args['colors_path']).astype('uint8') test(test_loader, test_data.data_list, net, 21, mean, std, 512, 480, 480, args['scales'], gray_folder, color_folder, colors)
def _create_data_loader(img_path, batch_size, img_size, n_cpu): """Creates a DataLoader for inferencing. :param img_path: Path to file containing all paths to validation images. :type img_path: str :param batch_size: Size of each image batch :type batch_size: int :param img_size: Size of each image dimension for yolo :type img_size: int :param n_cpu: Number of cpu threads to use during batch generation :type n_cpu: int :return: Returns DataLoader :rtype: DataLoader """ dataset = ImageFolder(img_path, transform=transforms.Compose( [DEFAULT_TRANSFORMS, Resize(img_size)])) dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=n_cpu, pin_memory=True) return dataloader
#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), batch_size=args.mini_batch_size, shuffle=True, **kwargs)
def main(): global BEST_LOSS cudnn.benchmark = True start_epoch = cfg.OPOSE.start_epoch # start from epoch 0 or last checkpoint epoch # Create ckpt & vis folder if not os.path.isdir(cfg.OPOSE.ckpt): os.makedirs(cfg.OPOSE.ckpt) if not os.path.exists(os.path.join(cfg.OPOSE.ckpt, 'vis')): os.makedirs(os.path.join(cfg.OPOSE.ckpt, 'vis')) if args.cfg_file is not None and not cfg.OPOSE.evaluate: shutil.copyfile( args.cfg_file, os.path.join(cfg.OPOSE.ckpt, args.cfg_file.split('/')[-1])) model = pose_estimation.PoseModel(num_point=19, num_vector=19, pretrained=True) # # Calculate FLOPs & Param # n_flops, n_convops, n_params = measure_model(model, cfg.OPOSE.input_size, cfg.OPOSE.input_size) criterion = nn.MSELoss().cuda() # Dataset and Loader train_dataset = dataset.CocoOpenposeData( cfg, cfg.OPOSE.data_root, cfg.OPOSE.info_root, 'train2017', transformer=transforms.Compose([ transforms.RandomResized(), transforms.RandomRotate(40), transforms.RandomCrop(368), transforms.RandomHorizontalFlip(), ])) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=int(cfg.workers), pin_memory=True) if cfg.OPOSE.validate or cfg.OPOSE.evaluate: val_dataset = dataset.CocoOpenposeData( cfg, cfg.OPOSE.data_root, cfg.OPOSE.info_root, 'val2017', transformer=transforms.Compose( [transforms.TestResized(cfg.OPOSE.input_size)])) val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=int(cfg.workers), pin_memory=True) # Load nets into gpu if NUM_GPUS > 1: model = torch.nn.DataParallel(model, device_ids=gpu).cuda() # Set up optimizers params, multiple = get_parameters(model, cfg, False) optimizer = torch.optim.SGD(params, cfg.OPOSE.base_lr, momentum=cfg.OPOSE.momentum, weight_decay=cfg.OPOSE.weight_decay) # Resume training title = 'Pytorch-OPOSE-{}-{}'.format(cfg.OPOSE.arch_encoder, cfg.OPOSE.arch_decoder) if cfg.OPOSE.resume: # Load checkpoint. print("==> Resuming from checkpoint '{}'".format(cfg.OPOSE.resume)) assert os.path.isfile( cfg.OPOSE.resume), 'Error: no checkpoint directory found!' ckpt = torch.load(cfg.OPOSE.resume) BEST_LOSS = ckpt['best_loss'] start_epoch = ckpt['epoch'] try: model.module.load_state_dict(ckpt['state_dict']) except: model.load_state_dict(ckpt['state_dict']) optimizer.load_state_dict(ckpt['optimizer']) logger = Logger(os.path.join(cfg.OPOSE.ckpt, 'log.txt'), title=title, resume=True) else: logger = Logger(os.path.join(cfg.OPOSE.ckpt, 'log.txt'), title=title) logger.set_names( ['epoch', 'Learning Rate', 'Train Loss', 'Valid Loss']) # Train and val for epoch in range(start_epoch, cfg.OPOSE.epochs): print('\nEpoch: [{}/{}] | LR: {:.8f} '.format(epoch + 1, cfg.OPOSE.epochs, cfg.OPOSE.base_lr)) train_loss = train(train_loader, model, criterion, optimizer, epoch, USE_CUDA) if cfg.OPOSE.validate: test_loss = test(val_loader, model, criterion, optimizer, epoch, USE_CUDA) else: test_loss = 0.0, 0.0 # Append logger file logger.append([epoch, cfg.OPOSE.base_lr, train_loss, test_loss]) # Save model save_checkpoint(model, optimizer, test_loss, epoch) # Adjust learning rate adjust_learning_rate(optimizer, epoch) # Draw curve try: draw_curve('model', cfg.OPOSE.ckpt) print('==> Success saving log curve...') except: print('==> Saving log curve error...') logger.close() try: savefig(os.path.join(cfg.OPOSE.ckpt, 'log.eps')) shutil.copyfile( os.path.join(cfg.OPOSE.ckpt, 'log.txt'), os.path.join( cfg.OPOSE.ckpt, 'log{}.txt'.format( datetime.datetime.now().strftime('%Y%m%d%H%M%S')))) except: print('Copy log error.') print('==> Training Done!') print('==> Best acc: {:.4f}%'.format(BEST_LOSS))