| # 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 torch | |
| from .task import Task | |
| class VLMTask(Task): | |
| """A VLM task for reproducibility. | |
| the collator split subsamples into two sub-batches. | |
| This has should have no logic changes. | |
| but changed the randomness in frame masking. | |
| """ | |
| def flat_subsample(self, tensor): | |
| size = tensor.size() | |
| if len(size) >= 2: | |
| batch_size = size[0] * (size[1] // 2) | |
| expanded_size = ( | |
| (batch_size, 2) + size[2:] if len(size) > 2 | |
| else (batch_size, 2) | |
| ) | |
| tensor = tensor.view(expanded_size) | |
| tensor = torch.cat([tensor[:, 0], tensor[:, 1]], dim=0) | |
| return tensor | |