Exemplo n.º 1
0
Arquivo: train.py Projeto: HCShi/IONet
def main():
    args = get_args()
    torch.backends.cudnn.enabled = False
    cudnn.benchmark = False
    torch.multiprocessing.set_sharing_strategy('file_system')

    train_loader = torch.utils.data.DataLoader(TrainDataset(
        args=args,
        transform=transforms.Compose([
            transforms.CenterCrop((224, 224)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor()
        ])),
                                               batch_size=args.batch_size,
                                               shuffle=True,
                                               num_workers=0)

    video_val_loader = torch.utils.data.DataLoader(VideoDataset(
        args=args,
        transform=transforms.Compose([
            transforms.CenterCrop((224, 224)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor()
        ]),
        test_mode=True),
                                                   batch_size=args.batch_size,
                                                   shuffle=False,
                                                   num_workers=0)

    print("start training")
    for epoch in range(args.epochs):
        train(train_loader, video_val_loader, args)
def get_checkpoint_loader() -> Any:
    """
    Get dataloder for extracting checkpoints.

    Returns
    -------
    loader : torch.utils.data.DataLoader
        Dataloader for extracting checkpoints.

    """
    filename = "./data/6zXXZvVvTFs"
    trim = (960, 1403)
    crop = (35, 50, 445, 300)
    dataset = VideoDataset(filename, trim, crop, frame_rate=15)
    loader = DataLoader(dataset, batch_size=32, shuffle=False, num_workers=1)

    return loader
Exemplo n.º 3
0
    def __init__(self, model,dataset_index=0,video_target = None):

        if args.video == None:
            
            self.video_target = video_target
            customset_train = CustomDataset(path = args.dataset_path,subset_type="training",dataset_index=dataset_index,video_target = video_target)
            customset_test = CustomDataset(path = args.dataset_path,subset_type="testing",dataset_index=dataset_index, video_target = video_target)
        
            self.trainloader = torch.utils.data.DataLoader(dataset=customset_train,batch_size=args.batch_size,shuffle=True,num_workers=args.num_workers)
            self.testloader = torch.utils.data.DataLoader(dataset=customset_test,batch_size=args.batch_size,shuffle=False,num_workers=args.num_workers)    
        else:
            video_dataset = VideoDataset(video=args.video, batch_size=args.batch_size,
                                        frame_skip=int(args.frame_skip),image_folder=args.extract_frames_path, use_existing=args.use_existing_frames)
            
            self.videoloader = torch.utils.data.DataLoader(dataset=video_dataset, batch_size=1,shuffle=False,num_workers=args.num_workers)

   
        if (model == "alex"):
            self.model = AlexNet()
        elif (model == "vgg"):
            self.model = VGG()
        elif (model == "resnet"):
            self.model = ResNet()

        if args.pretrained_model != None:
            if args.pretrained_finetuning == False:
                self.model.load_state_dict(torch.load(args.pretrained_model))
            else:
                print "DEBUG : Make it load only part of the resnet model"
                #print(self.model)
                #self.model.load_state_dict(torch.load(args.pretrained_model))
                #for param in self.model.parameters():
                #    param.requires_grad = False
                self.model.fc = nn.Linear(512, 1000)
                #print(self.model)
                self.model.load_state_dict(torch.load(args.pretrained_model))
                self.model.fc = nn.Linear(512,3)
                #print(self.model)
                
        self.model.cuda()        
        print "Using weight decay: ",args.weight_decay
        self.optimizer = optim.SGD(self.model.parameters(), weight_decay=float(args.weight_decay),lr=0.01, momentum=0.9,nesterov=True)
        self.criterion = nn.CrossEntropyLoss().cuda()
Exemplo n.º 4
0
def main():
    best_pec1 = 0
    args = get_args()

    torch.backends.cudnn.enabled = False
    cudnn.benchmark = False
    torch.multiprocessing.set_sharing_strategy('file_system')

    video_val_loader = torch.utils.data.DataLoader(VideoDataset(
        args=args,
        transform=transforms.Compose([
            transforms.CenterCrop((224, 224)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor()
        ]),
        test_mode=True),
                                                   batch_size=args.batch_size,
                                                   shuffle=True,
                                                   num_workers=8)

    print("start validate")
    validate(video_val_loader, args)
Exemplo n.º 5
0
    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)

use_gpu = torch.cuda.is_available()

######################################################################
def main(args):
    assert torch.cuda.is_available(), 'CUDA is not available.'
    torch.backends.cudnn.enabled = True
    torch.backends.cudnn.benchmark = True
    prepare_seed(args.rand_seed)

    logstr = 'seed-{:}-time-{:}'.format(args.rand_seed, time_for_file())
    logger = Logger(args.save_path, logstr)
    logger.log('Main Function with logger : {:}'.format(logger))
    logger.log('Arguments : -------------------------------')
    for name, value in args._get_kwargs():
        logger.log('{:16} : {:}'.format(name, value))
    logger.log("Python  version : {}".format(sys.version.replace('\n', ' ')))
    logger.log("Pillow  version : {}".format(PIL.__version__))
    logger.log("PyTorch version : {}".format(torch.__version__))
    logger.log("cuDNN   version : {}".format(torch.backends.cudnn.version()))

    # General Data Argumentation
    mean_fill = tuple([int(x * 255) for x in [0.485, 0.456, 0.406]])
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
    assert args.arg_flip == False, 'The flip is : {}, rotate is {}'.format(
        args.arg_flip, args.rotate_max)
    train_transform = [transforms.PreCrop(args.pre_crop_expand)]
    train_transform += [
        transforms.TrainScale2WH((args.crop_width, args.crop_height))
    ]
    train_transform += [
        transforms.AugScale(args.scale_prob, args.scale_min, args.scale_max)
    ]
    #if args.arg_flip:
    #  train_transform += [transforms.AugHorizontalFlip()]
    if args.rotate_max:
        train_transform += [transforms.AugRotate(args.rotate_max)]
    train_transform += [
        transforms.AugCrop(args.crop_width, args.crop_height,
                           args.crop_perturb_max, mean_fill)
    ]
    train_transform += [transforms.ToTensor(), normalize]
    train_transform = transforms.Compose(train_transform)

    eval_transform = transforms.Compose([
        transforms.PreCrop(args.pre_crop_expand),
        transforms.TrainScale2WH((args.crop_width, args.crop_height)),
        transforms.ToTensor(), normalize
    ])
    assert (
        args.scale_min + args.scale_max
    ) / 2 == args.scale_eval, 'The scale is not ok : {},{} vs {}'.format(
        args.scale_min, args.scale_max, args.scale_eval)

    # Model Configure Load
    model_config = load_configure(args.model_config, logger)
    args.sigma = args.sigma * args.scale_eval
    logger.log('Real Sigma : {:}'.format(args.sigma))

    # Training Dataset
    train_data = VDataset(train_transform, args.sigma, model_config.downsample,
                          args.heatmap_type, args.data_indicator,
                          args.video_parser)
    train_data.load_list(args.train_lists, args.num_pts, True)
    train_loader = torch.utils.data.DataLoader(train_data,
                                               batch_size=args.batch_size,
                                               shuffle=True,
                                               num_workers=args.workers,
                                               pin_memory=True)

    # Evaluation Dataloader
    eval_loaders = []
    if args.eval_vlists is not None:
        for eval_vlist in args.eval_vlists:
            eval_vdata = IDataset(eval_transform, args.sigma,
                                  model_config.downsample, args.heatmap_type,
                                  args.data_indicator)
            eval_vdata.load_list(eval_vlist, args.num_pts, True)
            eval_vloader = torch.utils.data.DataLoader(
                eval_vdata,
                batch_size=args.batch_size,
                shuffle=False,
                num_workers=args.workers,
                pin_memory=True)
            eval_loaders.append((eval_vloader, True))

    if args.eval_ilists is not None:
        for eval_ilist in args.eval_ilists:
            eval_idata = IDataset(eval_transform, args.sigma,
                                  model_config.downsample, args.heatmap_type,
                                  args.data_indicator)
            eval_idata.load_list(eval_ilist, args.num_pts, True)
            eval_iloader = torch.utils.data.DataLoader(
                eval_idata,
                batch_size=args.batch_size,
                shuffle=False,
                num_workers=args.workers,
                pin_memory=True)
            eval_loaders.append((eval_iloader, False))

    # Define network
    lk_config = load_configure(args.lk_config, logger)
    logger.log('model configure : {:}'.format(model_config))
    logger.log('LK configure : {:}'.format(lk_config))
    net = obtain_model(model_config, lk_config, args.num_pts + 1)
    assert model_config.downsample == net.downsample, 'downsample is not correct : {} vs {}'.format(
        model_config.downsample, net.downsample)
    logger.log("=> network :\n {}".format(net))

    logger.log('Training-data : {:}'.format(train_data))
    for i, eval_loader in enumerate(eval_loaders):
        eval_loader, is_video = eval_loader
        logger.log('The [{:2d}/{:2d}]-th testing-data [{:}] = {:}'.format(
            i, len(eval_loaders), 'video' if is_video else 'image',
            eval_loader.dataset))

    logger.log('arguments : {:}'.format(args))

    opt_config = load_configure(args.opt_config, logger)

    if hasattr(net, 'specify_parameter'):
        net_param_dict = net.specify_parameter(opt_config.LR, opt_config.Decay)
    else:
        net_param_dict = net.parameters()

    optimizer, scheduler, criterion = obtain_optimizer(net_param_dict,
                                                       opt_config, logger)
    logger.log('criterion : {:}'.format(criterion))
    net, criterion = net.cuda(), criterion.cuda()
    net = torch.nn.DataParallel(net)

    last_info = logger.last_info()
    if last_info.exists():
        logger.log("=> loading checkpoint of the last-info '{:}' start".format(
            last_info))
        last_info = torch.load(last_info)
        start_epoch = last_info['epoch'] + 1
        checkpoint = torch.load(last_info['last_checkpoint'])
        assert last_info['epoch'] == checkpoint[
            'epoch'], 'Last-Info is not right {:} vs {:}'.format(
                last_info, checkpoint['epoch'])
        net.load_state_dict(checkpoint['state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        scheduler.load_state_dict(checkpoint['scheduler'])
        logger.log("=> load-ok checkpoint '{:}' (epoch {:}) done".format(
            logger.last_info(), checkpoint['epoch']))
    elif args.init_model is not None:
        init_model = Path(args.init_model)
        assert init_model.exists(), 'init-model {:} does not exist'.format(
            init_model)
        checkpoint = torch.load(init_model)
        checkpoint = remove_module_dict(checkpoint['state_dict'], True)
        net.module.detector.load_state_dict(checkpoint)
        logger.log("=> initialize the detector : {:}".format(init_model))
        start_epoch = 0
    else:
        logger.log("=> do not find the last-info file : {:}".format(last_info))
        start_epoch = 0

    detector = torch.nn.DataParallel(net.module.detector)

    eval_results = eval_all(args, eval_loaders, detector, criterion,
                            'start-eval', logger, opt_config)
    if args.eval_once:
        logger.log("=> only evaluate the model once")
        logger.close()
        return

    # Main Training and Evaluation Loop
    start_time = time.time()
    epoch_time = AverageMeter()
    for epoch in range(start_epoch, opt_config.epochs):

        scheduler.step()
        need_time = convert_secs2time(
            epoch_time.avg * (opt_config.epochs - epoch), True)
        epoch_str = 'epoch-{:03d}-{:03d}'.format(epoch, opt_config.epochs)
        LRs = scheduler.get_lr()
        logger.log(
            '\n==>>{:s} [{:s}], [{:s}], LR : [{:.5f} ~ {:.5f}], Config : {:}'.
            format(time_string(), epoch_str, need_time, min(LRs), max(LRs),
                   opt_config))

        # train for one epoch
        train_loss = train(args, train_loader, net, criterion, optimizer,
                           epoch_str, logger, opt_config, lk_config,
                           epoch >= lk_config.start)
        # log the results
        logger.log('==>>{:s} Train [{:}] Average Loss = {:.6f}'.format(
            time_string(), epoch_str, train_loss))

        # remember best prec@1 and save checkpoint
        save_path = save_checkpoint(
            {
                'epoch': epoch,
                'args': deepcopy(args),
                'arch': model_config.arch,
                'state_dict': net.state_dict(),
                'detector': detector.state_dict(),
                'scheduler': scheduler.state_dict(),
                'optimizer': optimizer.state_dict(),
            },
            logger.path('model') /
            '{:}-{:}.pth'.format(model_config.arch, epoch_str), logger)

        last_info = save_checkpoint(
            {
                'epoch': epoch,
                'last_checkpoint': save_path,
            }, logger.last_info(), logger)

        eval_results = eval_all(args, eval_loaders, detector, criterion,
                                epoch_str, logger, opt_config)

        # measure elapsed time
        epoch_time.update(time.time() - start_time)
        start_time = time.time()

    logger.close()
Exemplo n.º 7
0
    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, 4, 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)

use_gpu = torch.cuda.is_available()


######################################################################
# Training the model
# --------
Exemplo n.º 8
0
######################################################################
# Load Data
# ---------
#

data_transforms = transforms.Compose([
    transforms.Resize((384, 128), interpolation=3),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

dataset = init_dataset('mars', root='../')

queryloader = DataLoader(
    VideoDataset(dataset.query, seq_len=opt.seq_len, sample_method=opt.sample_method,
    transform=data_transforms), batch_size=opt.batchsize, shuffle=False, num_workers=8, drop_last=False)

galleryloader = DataLoader(
    VideoDataset(dataset.gallery, seq_len=opt.seq_len,sample_method=opt.sample_method,
    transform=data_transforms), batch_size=opt.batchsize, shuffle=False, num_workers=8, drop_last=False)

use_gpu = torch.cuda.is_available()


######################################################################
# Load model
# --------
#

def load_network(network):
    save_path = os.path.join('./logs', name, 'net_%s.pth' % opt.which_epoch)
Exemplo n.º 9
0
def main():
    parser = argparse.ArgumentParser(description='VideoGAN')
    parser.add_argument('--batchsize', '-b', type=int, default=64,
                        help='Number of videos in each mini-batch')
    parser.add_argument('--epoch', '-e', type=int, default=1000,
                        help='Number of sweeps over the dataset to train')
    parser.add_argument('--gpu', '-g', type=int, default=-1,
                        help='GPU ID (negative value indicates CPU)')
    parser.add_argument('--dataset', '-i', required=True,
                        help='Directory list of images files')
    parser.add_argument('--root', default='.',
                        help='Path of images files')
    parser.add_argument('--out', '-o', default='result',
                        help='Directory to output the result')
    parser.add_argument('--resume', '-r', default='',
                        help='Resume the training from snapshot')
    parser.add_argument('--seed', type=int, default=0,
                        help='Random seed of z at visualization stage')
    parser.add_argument('--snapshot_interval', type=int, default=1000,
                        help='Interval of snapshot')
    parser.add_argument('--display_interval', type=int, default=100,
                        help='Interval of displaying log to console')
    parser.add_argument('--video_codecs', default='H264',
                        help='Video Codec')
    parser.add_argument('--video_ext', default='avi',
                        help='Extension of output video files')
    args = parser.parse_args()

    print('GPU: {}'.format(args.gpu))
    print('# Minibatch-size: {}'.format(args.batchsize))
    print('# epoch: {}'.format(args.epoch))

    gen = Generator()
    dis = Discriminator()

    if args.gpu >= 0:
        chainer.cuda.get_device_from_id(args.gpu).use()
        gen.to_gpu()
        dis.to_gpu()

    def make_optimizer(model, alpha=0.0002, beta1=0.5):
        optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1)
        optimizer.setup(model)
        optimizer.add_hook(chainer.optimizer.WeightDecay(0.0001), 'hook_dec')
        return optimizer
    opt_gen = make_optimizer(gen)
    opt_dis = make_optimizer(dis)

    train = VideoDataset(paths=_get_images_paths(args.dataset, args.root))
    print('# data-size: {}'.format(len(train)))
    print('# data-shape: {}'.format(train[0].shape))
    print('')

    train_iter = chainer.iterators.SerialIterator(train, args.batchsize)

    updater = VideoGANUpdater(
        models=(gen, dis),
        iterator=train_iter,
        optimizer={'gen': opt_gen, 'dis': opt_dis},
        device=args.gpu)
    trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)

    snapshot_interval = (args.snapshot_interval, 'iteration')
    display_interval = (args.display_interval, 'iteration')
    trainer.extend(
        extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'),
        trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        gen, 'gen_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
    trainer.extend(extensions.snapshot_object(
        dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
    trainer.extend(extensions.LogReport(trigger=display_interval))
    trainer.extend(extensions.PrintReport([
        'epoch', 'iteration', 'gen/loss', 'dis/loss',
    ]), trigger=display_interval)
    trainer.extend(extensions.ProgressBar(update_interval=10))
    trainer.extend(
        out_generated_video(
            gen, dis, 5, args.seed, args.out,
            args.video_codecs, args.video_ext),
        trigger=snapshot_interval)

    if args.resume:
        chainer.serializers.load_npz(args.resume, trainer)

    trainer.run()

    if args.gpu >= 0:
        gen.to_cpu()
    chainer.serializers.save_npz(os.path.join(args.out, 'gen.npz'), gen)
Exemplo n.º 10
0
classez = ['Archery', 'BalanceBeam', 'BaseballPitch', 'BenchPress', 'Biking',
'Bowling', 'Fencing', 'HammerThrow', 'HighJump', 'JavelinThrow', 'Kayaking',
'LongJump', 'PoleVault', 'PommelHorse', 'Rowing', 'SkateBoarding', 'SkyDiving',
'SumoWrestling', 'Surfing', 'TrampolineJumping']

# Init Network and optiization parameters
# You can tweak this parameters but that is not the homework
net = resnet_tsn(pretrained=True, progress=True)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=40, gamma=0.1)

# Set up datasets and pytorch loaders, adjust batchsize for your GPU, now it uses about 2.1 GB
# Tweak CPU usage here
train_set = VideoDataset('/home/jcleon/tmp/train', transform_train, classez)
val_set = VideoDatasetVal('/home/jcleon/tmp/val', transform_val, classez)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=24, shuffle=True,
                                          num_workers=4, pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_set, batch_size=24, shuffle=False,
                                          num_workers=4, pin_memory=True)

#GPU selection, if single gpu use cuda:0 instead of cuda:3
has_cuda = torch.cuda.is_available()
device = torch.device('cuda:3' if has_cuda else 'cpu')
net = net.to(device)

# Training loop
for epoch in range(100):
    scheduler.step()
    print('Epoch ', epoch+1)
Exemplo n.º 11
0
            for version in args.version_valid
        ]
    else:
        train_files += [
            os.path.join(args.csv_location,
                         '%s-%s-train.csv' % (d, args.version))
        ]
        valid_files += [
            os.path.join(args.csv_location,
                         '%s-%s-validation.csv' % (d, args.version))
        ]

dataset = dict(train=VideoDataset(train_files,
                                  labels,
                                  select='random',
                                  transform=transform_train,
                                  sample_length=args.sample_length,
                                  sample_size=args.sample_size,
                                  location=args.location),
               valid=VideoDataset(valid_files,
                                  labels,
                                  select='center',
                                  transform=transform_valid,
                                  sample_length=args.sample_length,
                                  sample_size=args.sample_size,
                                  location=args.test_location))

sampler = EqualDistSampler(dataset['train']._targets)
train_loader = DataLoader(dataset["train"],
                          batch_size=args.batch_size,
                          num_workers=args.num_workers,