def init_dataset(use_gpu): normalize = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) transform_tr = Compose([ Random2DTranslation(args.height, args.width, p=0.5), RandomHorizontalFlip(), ToTensor(), normalize ]) transform_te = Compose( [Resize([args.height, args.width]), ToTensor(), normalize]) trainset = datasets.init_dataset(args.dataset, root=args.root, transform=transform_tr, mode='train', verbose=True) valset = datasets.init_dataset(args.dataset, root=args.root, transform=transform_te, mode='val', verbose=False) testset = datasets.init_dataset(args.dataset, root=args.root, transform=transform_te, mode='test', verbose=False) num_attrs = trainset.num_attrs attr_dict = trainset.attr_dict trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=use_gpu, drop_last=True) valloader = torch.utils.data.DataLoader(valset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=use_gpu, drop_last=False) testloader = torch.utils.data.DataLoader(testset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=use_gpu, drop_last=False) return trainloader, valloader, testloader, num_attrs, attr_dict
def run(): args = parse_opts() os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" # see issue #152 os.environ["CUDA_VISIBLE_DEVICES"] = "9" # GLOBAL VARS # MODE = args.mode CLASS_WEIGHT = False N_EP = 20 FLATTEN = args.flatten RNN = args.rnn BATCH_SIZE = args.batch_size #### datasets, dataloaders = init_dataset( BATCH_SIZE, single_channel=args.single_channel) print('[Train] class counts', np.unique( datasets['train'].target_vals, return_counts=True)) print('[Test] class counts', np.unique( datasets['test'].target_vals, return_counts=True)) n_ch = 1 if args.single_channel else 3 if MODE == 'min': in_channels = datasets['train'].min_depth*n_ch elif MODE == 'max': in_channels = datasets['train'].max_depth*n_ch torch.manual_seed(0) # init net net = init_net(opt=args.model_idx, in_channels=in_channels) class_weight = None if CLASS_WEIGHT: cnts = Counter(datasets['train'].target_vals) n = len(datasets['train']) class_weight = [max(cnts.values())/cnts['0'], max(cnts.values())/cnts['1']] class_weight = torch.FloatTensor(class_weight) cross_entrp_loss = nn.CrossEntropyLoss(weight=class_weight).cuda() focal_loss = FocalLoss().cuda() optimizer = optim.Adam(net.parameters(), lr=0.000027) criterion = cross_entrp_loss # scheduler = optim.lr_scheduler.ReduceLROnPlateau( # optimizer, 'min', verbose=True, patience=7) for ep in range(N_EP): train_epoch(net, dataloaders['train'], optimizer, criterion, ep, scheduler=None, flatten=FLATTEN, MODE=MODE, rnn=RNN) valid_loss = evaluate(net, dataloaders['test'], criterion, ep, flatten=FLATTEN, MODE=MODE, rnn=RNN)
def main(): global args args = parser.parse_args() os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_ids cudnn.benchmark = True args.resume = osp.join(args.resume, 'best_model.pth.tar') save_root = osp.join(osp.dirname(args.resume), 'results') if not osp.isdir(save_root): os.makedirs(save_root) print("==========\nArgs:{}\n==========".format(args)) # create model print("=> creating model '{}'".format(args.arch)) model = models.init_model(name=args.arch) print("Model size: {:.5f}M".format(sum(p.numel() for p in model.parameters())/1000000.0)) if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume) model.load_state_dict(checkpoint['state_dict']) else: print("=> no checkpoint found at '{}'".format(args.resume)) return model = torch.nn.DataParallel(model).cuda() print("Initializing dataset {}".format(args.dataset)) dataset = datasets.init_dataset(args.dataset, root=osp.join(args.data_root,args.dataset)) to_pil = T.ToPILImage() model.eval() print("===> Start testing") with torch.no_grad(): for img in dataset.valset_select['raw']: raw_path = osp.join(args.data_root,args.dataset,img) raw_pil = Image.open(raw_path) raw = depth_transform(raw_pil) raw = TF.to_tensor(raw).float() valid_mask = (raw>0).detach().float() input = torch.unsqueeze(raw,0).cuda() output = torch.clamp(model(input),min=0, max=255) output = output * 256. raw = raw * 256. output = output[0].cpu() output = raw*valid_mask + output*(1-valid_mask) pil_img = to_pil(output.int()) pil_img.save(osp.join(save_root, osp.basename(img))) print(img+' finish.')
train_transforms = T.Compose([ T.Resize([384, 128]), T.RandomHorizontalFlip(p=0.5), T.Pad(10), T.RandomCrop([384, 128]), T.ToTensor(), normalize_transform, RandomErasing(probability=0.5, mean=[0.485, 0.456, 0.406]) ]) # val_transforms = T.Compose([ # T.Resize([384, 128]), # T.ToTensor(), # normalize_transform # ]) dataset = init_dataset('mars', root='../') dataset_sizes = {} dataset_sizes['train'] = dataset.num_train_imgs train_set = VideoDataset(dataset.train, opt.seq_len, opt.sample_method, train_transforms) dataloaders = {} dataloaders['train'] = DataLoader(train_set, batch_size=opt.batchsize, drop_last=True, sampler=RandomIdentitySampler( dataset.train, opt.batchsize, 4), num_workers=8) # val_set = VideoDataset(dataset.query + dataset.gallery, 4, val_transforms) # dataloaders['val'] = DataLoader( # val_set, batch_size=opt.batchsize, drop_last=True, shuffle=False, num_workers=8)
class_names = image_datasets['train'].classes use_gpu = torch.cuda.is_available() # since = time.time() # inputs, classes = next(iter(dataloaders['train'])) # print(time.time()-since) ###################################################################### # New Train Loader # -------- # dataset = init_dataset('market1501', root='../') train_set = ImageDataset(dataset.train, data_transforms['train']) dataloaders['train'] = DataLoader( train_set, batch_size=opt.batchsize, drop_last=True, sampler=RandomIdentitySampler(dataset.train, opt.batchsize, 4), num_workers=8) ###################################################################### # Training the model # -------- # y_loss = {} # loss history y_loss['train'] = [] y_loss['val'] = []
def init_dataset(self): """ """ self.dataset = init_dataset(**self.params['dataset']) self.params['model']['num_classes'] = self.dataset.num_classes