Spaces:
Sleeping
Sleeping
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| # | |
| # This source code is licensed under the MIT license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| import os | |
| import torch | |
| import pickle | |
| import random | |
| from tqdm import tqdm | |
| from torch.utils.data import DataLoader | |
| from torch.utils.data.distributed import DistributedSampler | |
| from ..processors import ( | |
| ShardedHow2MetaProcessor, | |
| ShardedVideoProcessor, | |
| ShardedTextProcessor, | |
| VariedLenAligner, | |
| ) | |
| from ..datasets import MMDataset | |
| from .task import Task | |
| from ..modules import vectorpool | |
| from ..evaluators.predictor import Predictor | |
| from ..utils import set_seed, get_local_rank, get_world_size | |
| class RetriTask(Task): | |
| """abstract class for task with retrival.""" | |
| def reshape_subsample(self, sample): | |
| for key in sample: | |
| if torch.is_tensor(sample[key]): | |
| sample[key] = self.flat_subsample(sample[key]) | |
| return sample | |
| def flat_subsample(self, tensor): | |
| if tensor.size(0) == 1: | |
| tensor = tensor.squeeze(0) | |
| return tensor | |
| def build_dataloader(self): | |
| """called by `get_batch_iterator` in fairseqmmtask. """ | |
| # TODO: hard-code dataloader for retri for now and configurable in .yaml. | |
| # reuse the `train.lst`. | |
| self.config.dataset.split = "train" | |
| meta_processor = ShardedHow2MetaProcessor(self.config.dataset) | |
| video_processor = ShardedVideoProcessor(self.config.dataset) | |
| text_processor = ShardedTextProcessor(self.config.dataset) | |
| aligner = VariedLenAligner(self.config.dataset) | |
| aligner.subsampling = self.config.dataset.clip_per_video | |
| self.retri_data = MMDataset( | |
| meta_processor, video_processor, text_processor, aligner | |
| ) | |
| retri_sampler = DistributedSampler(self.retri_data) | |
| infer_scale = 16 | |
| batch_size = self.config.dataset.num_video_per_batch \ | |
| * infer_scale | |
| self.retri_dataloader = DataLoader( | |
| self.retri_data, | |
| collate_fn=self.retri_data.collater, | |
| batch_size=batch_size, | |
| shuffle=False, | |
| sampler=retri_sampler, | |
| num_workers=self.config.fairseq.dataset.num_workers | |
| ) | |
| return self.retri_dataloader | |
| def retrive_candidates(self, epoch, dataloader=None): | |
| if get_local_rank() == 0: | |
| print("running retrieval model.") | |
| out_dir = os.path.join( | |
| self.config.fairseq.checkpoint.save_dir, "retri") | |
| os.makedirs(out_dir, exist_ok=True) | |
| if not os.path.isfile( | |
| os.path.join( | |
| out_dir, "batched_e" + str(epoch) + "_videos0.pkl") | |
| ): | |
| if dataloader is None: | |
| dataloader = self.retri_dataloader | |
| self.model.eval() | |
| self.model.is_train = False | |
| assert self.retri_data.meta_processor.data == \ | |
| self.train_data.meta_processor.data # video_ids not mutated. | |
| self._retri_predict(epoch, dataloader) | |
| self.model.train() | |
| self.model.is_train = True | |
| torch.distributed.barrier() | |
| output = self._retri_sync(epoch, out_dir) | |
| torch.distributed.barrier() | |
| self.train_data.meta_processor.set_candidates(output) | |
| return output | |
| class VideoRetriTask(RetriTask): | |
| """RetriTask on video level.""" | |
| def reshape_subsample(self, sample): | |
| if ( | |
| hasattr(self.config.dataset, "clip_per_video") | |
| and self.config.dataset.clip_per_video is not None | |
| and self.config.dataset.clip_per_video > 1 | |
| ): | |
| for key in sample: | |
| if torch.is_tensor(sample[key]): | |
| sample[key] = self.flat_subsample(sample[key]) | |
| return sample | |
| def flat_subsample(self, tensor): | |
| if tensor.size(0) == 1: | |
| tensor = tensor.squeeze(0) | |
| return Task.flat_subsample(self, tensor) | |
| def _retri_predict(self, epoch, dataloader): | |
| set_seed(epoch) | |
| # save for retrival. | |
| predictor = VideoPredictor(self.config) | |
| predictor.predict_loop( | |
| self.model, dataloader) | |
| set_seed(epoch) # get the same text clips. | |
| # retrival. | |
| retri_predictor = VideoRetriPredictor( | |
| self.config) | |
| retri_predictor.predict_loop( | |
| self.model, predictor.vecpool.retriver, epoch) | |
| del predictor | |
| del retri_predictor | |
| def _retri_sync(self, epoch, out_dir): | |
| # gpu do the same merge. | |
| batched_videos = [] | |
| for local_rank in range(get_world_size()): | |
| fn = os.path.join( | |
| out_dir, | |
| "batched_e" + str(epoch) + "_videos" + str(local_rank) + ".pkl") | |
| with open(fn, "rb") as fr: | |
| batched_videos.extend(pickle.load(fr)) | |
| print( | |
| "[INFO] batched_videos", | |
| len(batched_videos), len(batched_videos[0])) | |
| return batched_videos | |
| class VideoPredictor(Predictor): | |
| def __init__(self, config): | |
| vectorpool_cls = getattr(vectorpool, config.vectorpool_cls) | |
| self.vecpool = vectorpool_cls(config) | |
| def predict_loop( | |
| self, | |
| model, | |
| dataloader, | |
| early_stop=-1, | |
| ): | |
| with torch.no_grad(): | |
| if get_local_rank() == 0: | |
| dataloader = tqdm(dataloader) | |
| for batch_idx, batch in enumerate(dataloader): | |
| if batch_idx == early_stop: | |
| break | |
| self(batch, model) | |
| return self.finalize() | |
| def __call__(self, sample, model, **kwargs): | |
| param = next(model.parameters()) | |
| dtype = param.dtype | |
| device = param.device | |
| subsample = sample["vfeats"].size(1) | |
| sample = self.to_ctx(sample, device, dtype) | |
| for key in sample: | |
| if torch.is_tensor(sample[key]): | |
| size = sample[key].size() | |
| if len(size) >= 2: | |
| batch_size = size[0] * size[1] | |
| expanded_size = ( | |
| (batch_size,) + size[2:] if len(size) > 2 | |
| else (batch_size,) | |
| ) | |
| sample[key] = sample[key].view(expanded_size) | |
| outputs = model(**sample) | |
| sample.update(outputs) | |
| self.vecpool(sample, subsample) | |
| def finalize(self): | |
| print("[INFO]", self.vecpool) | |
| if not self.vecpool.retriver.db.is_trained: | |
| self.vecpool.retriver.finalize_training() | |
| return self.vecpool.retriver | |
| class VideoRetriPredictor(Predictor): | |
| """ | |
| Online Retrieval Predictor for Clips (used by RetriTask). | |
| TODO: merge this with VisPredictor? | |
| """ | |
| def __init__(self, config): | |
| self.pred_dir = os.path.join( | |
| config.fairseq.checkpoint.save_dir, | |
| "retri") | |
| self.num_cands = config.num_cands | |
| self.num_video_per_batch = config.dataset.num_video_per_batch | |
| def predict_loop( | |
| self, | |
| model, | |
| retriver, | |
| epoch, | |
| early_stop=-1 | |
| ): | |
| # a fake loop that only try to recover video vector | |
| # from video_id. | |
| batched_videos = [] | |
| # obtain available video_ids. | |
| video_ids = list(retriver.videoid_to_vectoridx.keys()) | |
| dataloader = random.sample( | |
| video_ids, | |
| len(video_ids) // self.num_video_per_batch | |
| ) | |
| if get_local_rank() == 0: | |
| dataloader = tqdm(dataloader) | |
| for batch_idx, batch in enumerate(dataloader): | |
| # batch is one video id. | |
| if batch_idx == early_stop: | |
| break | |
| video_ids = retriver.search_by_video_ids( | |
| [batch], self.num_cands)[0] | |
| if len(video_ids) > self.num_video_per_batch: | |
| # we moved the center to make cluster robust. | |
| video_ids = random.sample(video_ids, self.num_video_per_batch) | |
| batched_videos.append(video_ids) | |
| return self.finalize(batched_videos, epoch) | |
| def finalize(self, batched_videos, epoch): | |
| fn = os.path.join( | |
| self.pred_dir, | |
| "batched_e" + str(epoch) + "_videos" + str(get_local_rank()) + ".pkl") | |
| with open(fn, "wb") as fw: | |
| pickle.dump(batched_videos, fw, pickle.HIGHEST_PROTOCOL) | |
| return batched_videos | |