File size: 1,839 Bytes
a3c8a6a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | import torch
from torch.utils.data import DataLoader
from dataloaders.dataloader_msvd_retrieval import MSVD_DataLoader
def dataloader_msvd_train(args, tokenizer):
msvd_dataset = MSVD_DataLoader(
subset="train",
data_path=args.data_path,
features_path=args.features_path,
csv_path=args.train_csv,
max_words=args.max_words,
feature_framerate=args.feature_framerate,
tokenizer=tokenizer,
max_frames=args.max_frames,
frame_order=args.train_frame_order,
slice_framepos=args.slice_framepos,
)
train_sampler = torch.utils.data.distributed.DistributedSampler(msvd_dataset)
dataloader = DataLoader(
msvd_dataset,
batch_size=args.batch_size // args.n_gpu,
num_workers=args.num_thread_reader,
pin_memory=False,
shuffle=(train_sampler is None),
sampler=train_sampler,
drop_last=True,
)
return dataloader, len(msvd_dataset), train_sampler
def dataloader_msvd_test(args, tokenizer, subset="test"):
msvd_testset = MSVD_DataLoader(
subset=subset,
data_path=args.data_path,
features_path=args.features_path,
csv_path=args.val_csv,
max_words=args.max_words,
feature_framerate=args.feature_framerate,
tokenizer=tokenizer,
max_frames=args.max_frames,
frame_order=args.eval_frame_order,
slice_framepos=args.slice_framepos,
)
dataloader_msrvtt = DataLoader(
msvd_testset,
batch_size=args.batch_size_val,
num_workers=args.num_thread_reader,
shuffle=False,
drop_last=False,
)
return dataloader_msrvtt, len(msvd_testset)
DATALOADER_DICT = {}
DATALOADER_DICT["msvd"] = {"train":dataloader_msvd_train, "val":dataloader_msvd_test, "test":dataloader_msvd_test}
|