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
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()
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)
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()
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
# --------
######################################################################
# 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)
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)
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)
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,