def __init__(self, args, device):
super(NTUSearcher, self).__init__(args)
self.device = device
# Handle data
transformer_val = transforms.Compose(
[ntu_data.NormalizeLen(args.vid_len),
ntu_data.ToTensor()])
transformer_tra = transforms.Compose([
ntu_data.AugCrop(),
ntu_data.NormalizeLen(args.vid_len),
ntu_data.ToTensor()
])
dataset_training = ntu_data.NTU(args.datadir,
transform=transformer_tra,
stage='trainexp',
args=args)
dataset_dev = ntu_data.NTU(args.datadir,
transform=transformer_val,
stage='dev',
args=args)
datasets = {'train': dataset_training, 'dev': dataset_dev}
self.dataloaders = {
x: DataLoader(datasets[x],
batch_size=args.batchsize,
shuffle=True,
num_workers=args.num_workers,
drop_last=False)
for x in ['train', 'dev']
}
def get_dataloaders(args):
import torchvision.transforms as transforms
from datasets import ntu as d
from torch.utils.data import DataLoader
# Handle data
#transformer_val = transforms.Compose([d.NormalizeLen(args.vid_len), d.ToTensor()])
#transformer_tra = transforms.Compose([d.AugCrop(), d.NormalizeLen(args.vid_len), d.ToTensor()])
transformer_val = transforms.Compose([
d.VisualRandomCrop(cropsize=(224, 224), central=True),
d.NormalizeLen(args.vid_len),
d.ToTensor()
])
transformer_tra = transforms.Compose([
d.VisualRandomCrop(cropsize=(224, 224), central=True),
d.AugCrop(),
d.NormalizeLen(args.vid_len),
d.ToTensor()
])
dataset_training = d.NTU(args.datadir,
transform=transformer_tra,
stage='train',
args=args)
dataset_testing = d.NTU(args.datadir,
transform=transformer_val,
stage='test',
args=args)
dataset_validation = d.NTU(args.datadir,
transform=transformer_val,
stage='dev',
args=args)
datasets = {
'train': dataset_training,
'dev': dataset_validation,
'test': dataset_testing
}
dataloaders = {
x: DataLoader(datasets[x],
batch_size=args.batchsize,
shuffle=True,
num_workers=args.num_workers,
drop_last=False,
pin_memory=True)
for x in ['train', 'dev', 'test']
}
return dataloaders