100 lines
4.2 KiB
Python
100 lines
4.2 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
|
|
# holder of all proprietary rights on this computer program.
|
|
# You can only use this computer program if you have closed
|
|
# a license agreement with MPG or you get the right to use the computer
|
|
# program from someone who is authorized to grant you that right.
|
|
# Any use of the computer program without a valid license is prohibited and
|
|
# liable to prosecution.
|
|
#
|
|
# Copyright©2020 Max-Planck-Gesellschaft zur Förderung
|
|
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
|
|
# for Intelligent Systems. All rights reserved.
|
|
#
|
|
# Contact: ps-license@tuebingen.mpg.de
|
|
|
|
from typing import List, Dict
|
|
from torch import Tensor
|
|
|
|
|
|
def collate_tensor_with_padding(batch: List[Tensor]) -> Tensor:
|
|
dims = batch[0].dim()
|
|
max_size = [max([b.size(i) for b in batch]) for i in range(dims)]
|
|
size = (len(batch),) + tuple(max_size)
|
|
canvas = batch[0].new_zeros(size=size)
|
|
for i, b in enumerate(batch):
|
|
sub_tensor = canvas[i]
|
|
for d in range(dims):
|
|
sub_tensor = sub_tensor.narrow(d, 0, b.size(d))
|
|
sub_tensor.add_(b)
|
|
return canvas
|
|
|
|
|
|
def collate_datastruct_and_text(lst_elements: List) -> Dict:
|
|
collate_datastruct = lst_elements[0]["datastruct"].transforms.collate
|
|
|
|
batch = {
|
|
# Collate with padding for the datastruct
|
|
"datastruct": collate_datastruct([x["datastruct"] for x in lst_elements]),
|
|
# Collate normally for the length
|
|
"length": [x["length"] for x in lst_elements],
|
|
# Collate the text
|
|
"text": [x["text"] for x in lst_elements]}
|
|
|
|
# add keyid for example
|
|
otherkeys = [x for x in lst_elements[0].keys() if x not in batch]
|
|
for key in otherkeys:
|
|
batch[key] = [x[key] for x in lst_elements]
|
|
|
|
return batch
|
|
|
|
def collate_length_and_text(lst_elements: List) -> Dict:
|
|
|
|
batch = {
|
|
"length_0": [x["length_0"] for x in lst_elements],
|
|
"length_1": [x["length_1"] for x in lst_elements],
|
|
"length_transition": [x["length_transition"] for x in lst_elements],
|
|
"length_1_with_transition": [x["length_1_with_transition"] for x in lst_elements],
|
|
"text_0": [x["text_0"] for x in lst_elements],
|
|
"text_1": [x["text_1"] for x in lst_elements]
|
|
}
|
|
|
|
return batch
|
|
|
|
def collate_pairs_and_text(lst_elements: List, ) -> Dict:
|
|
if 'features_0' not in lst_elements[0]: # test set
|
|
collate_datastruct = lst_elements[0]["datastruct"].transforms.collate
|
|
batch = {"datastruct": collate_datastruct([x["datastruct"] for x in lst_elements]),
|
|
"length_0": [x["length_0"] for x in lst_elements],
|
|
"length_1": [x["length_1"] for x in lst_elements],
|
|
"length_transition": [x["length_transition"] for x in lst_elements],
|
|
"length_1_with_transition": [x["length_1_with_transition"] for x in lst_elements],
|
|
"text_0": [x["text_0"] for x in lst_elements],
|
|
"text_1": [x["text_1"] for x in lst_elements]
|
|
}
|
|
|
|
else:
|
|
batch = {"motion_feats_0": collate_tensor_with_padding([el["features_0"] for el in lst_elements]),
|
|
"motion_feats_1": collate_tensor_with_padding([el["features_1"] for el in lst_elements]),
|
|
"motion_feats_1_with_transition": collate_tensor_with_padding([el["features_1_with_transition"] for el in lst_elements]),
|
|
"length_0": [x["length_0"] for x in lst_elements],
|
|
"length_1": [x["length_1"] for x in lst_elements],
|
|
"length_transition": [x["length_transition"] for x in lst_elements],
|
|
"length_1_with_transition": [x["length_1_with_transition"] for x in lst_elements],
|
|
"text_0": [x["text_0"] for x in lst_elements],
|
|
"text_1": [x["text_1"] for x in lst_elements]
|
|
}
|
|
return batch
|
|
|
|
|
|
def collate_text_and_length(lst_elements: Dict) -> Dict:
|
|
batch = {"length": [x["length"] for x in lst_elements],
|
|
"text": [x["text"] for x in lst_elements]}
|
|
|
|
# add keyid for example
|
|
otherkeys = [x for x in lst_elements[0].keys() if x not in batch and x != "datastruct"]
|
|
for key in otherkeys:
|
|
batch[key] = [x[key] for x in lst_elements]
|
|
return batch
|