log_dir = os.path.join(args.log, exp_name)
writer = SummaryWriter(log_dir)
train_transforms = transforms.Compose([
transforms.Resize((128, 171)),
transforms.RandomCrop(112),
transforms.ToTensor()
])
if args.dataset == 'ucf101':
train_dataset = UCF101Dataset('data/ucf101', args.cl, args.split, True, train_transforms)
# 交叉验证集选取,按比例
val_size = 800
train_dataset, val_dataset = random_split(train_dataset, (len(train_dataset) - val_size, val_size))
elif args.dataset == 'hmdb51':
train_dataset = HMDB51Dataset('data/hmdb51', args.cl, args.split, True, train_transforms)
val_size = 400
train_dataset, val_dataset = random_split(train_dataset, (len(train_dataset) - val_size, val_size))
elif args.dataset == "monkey2":
train_dataset = MONKEY2Dataset('data/monkey2', args.cl, args.split, True, train_transforms)
val_size = 20
train_dataset, val_dataset = random_split(train_dataset, (len(train_dataset) - val_size, val_size))
print('TRAIN video number: {}, VAL video number: {}.'.format(len(train_dataset), len(val_dataset)))
train_dataloader = DataLoader(train_dataset, batch_size=args.bs, shuffle=True,
num_workers=args.workers, pin_memory=True)
val_dataloader = DataLoader(val_dataset, batch_size=args.bs, shuffle=False,
num_workers=args.workers, pin_memory=True)
if args.ckpt:
else:
exp_name = 'K400_TCG_split1_finetuned_loss_{}_cl{}_{}'.format(args.model, args.cl, time.strftime('%m%d%H%M'))
log_dir = os.path.join(args.log, exp_name)
writer = SummaryWriter(log_dir)
train_transforms = transforms.Compose([
transforms.Resize((128, 171)),
transforms.RandomCrop(112),
transforms.ToTensor()
])
if args.dataset == 'ucf101':
train_dataset = UCF101Dataset('data/ucf101', args.cl, args.split, True, train_transforms)
val_size = 800
elif args.dataset == 'hmdb51':
train_dataset = HMDB51Dataset('data/hmdb51', args.cl, args.split, True, train_transforms)
val_size = 400
elif args.dataset == 'K400':
train_dataset = K400Dataset_train('data/K400', args.cl, args.split, True, train_transforms)
val_dataset = K400Dataset_val('data/K400', args.cl, args.split, True, train_transforms)
# split val for 800 videos
#train_dataset, val_dataset = random_split(train_dataset, (len(train_dataset)-val_size, val_size))
print('TRAIN video number: {}, VAL video number: {}.'.format(len(train_dataset), len(val_dataset)))
train_dataloader = DataLoader(train_dataset, batch_size=args.bs, shuffle=True,
num_workers=args.workers, pin_memory=True)
val_dataloader = DataLoader(val_dataset, batch_size=args.bs, shuffle=False,
num_workers=args.workers, pin_memory=True)
# save graph and clips_order samples
for data in train_dataloader:
def main():
if not torch.cuda.is_available():
raise 'Only support GPU mode'
# parse the args
args = parse_option()
print(vars(args))
best_acc = 0 # best test accuracy
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
print('[Warning] The training modalities are RGB and [{}]'.format(
args.modality))
# Data
train_transforms = transforms.Compose([
transforms.Resize((128, 171)), # smaller edge to 128
transforms.RandomCrop(112),
transforms.ToTensor()
])
if args.dataset == 'ucf101':
trainset = UCF101Dataset('./data/ucf101/',
transforms_=train_transforms)
else:
trainset = HMDB51Dataset('./data/hmdb51/',
transforms_=train_transforms)
train_loader = DataLoader(trainset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True)
n_data = trainset.__len__()
# set the model
model, contrast, criterion_1, criterion_2 = set_model(args, n_data)
# set the optimizer
optimizer = set_optimizer(args, model)
# optionally resume from a checkpoint
args.start_epoch = 1
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume, map_location='cpu')
args.start_epoch = checkpoint['epoch'] + 1
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
contrast.load_state_dict(checkpoint['contrast'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
del checkpoint
torch.cuda.empty_cache()
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# tensorboard
logger = tb_logger.Logger(logdir=args.tb_folder, flush_secs=2)
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=[45, 90, 125, 160],
gamma=0.2)
# routine
for epoch in range(args.start_epoch, args.epochs + 1):
time1 = time.time()
view1_loss, view1_prob, view2_loss, view2_prob = train(
epoch, train_loader, model, contrast, criterion_1, criterion_2,
optimizer, args)
time2 = time.time()
print('\nepoch {}, total time {:.2f}'.format(epoch, time2 - time1))
# tensorboard logger
logger.log_value('view1_loss', view1_loss, epoch)
logger.log_value('view1_prob', view1_prob, epoch)
logger.log_value('view2_loss', view2_loss, epoch)
logger.log_value('view2_prob', view2_prob, epoch)
# save model
if epoch % args.save_freq == 0:
print('==> Saving...')
state = {
'opt': args,
'model': model.state_dict(),
'contrast': contrast.state_dict(),
'optimizer': optimizer.state_dict(),
'epoch': epoch,
}
save_file = os.path.join(
args.model_folder,
'ckpt_epoch_{epoch}.pth'.format(epoch=epoch))
torch.save(state, save_file)
# help release GPU memory
del state
torch.cuda.empty_cache()
scheduler.step()
print(args.model_name)
