| from collections import defaultdict | |
| import numpy as np | |
| from torch.nn.utils.rnn import pad_sequence | |
| def channels_to_2nd_dim(x): | |
| return x.permute(*([0, x.dim() - 1] + list(range(1, x.dim() - 1)))) | |
| def channels_to_final_dim(x): | |
| return x.permute(*([0] + list(range(2, x.dim())) + [1])) | |
| def collate(tensor_list): | |
| out_dict = defaultdict() | |
| for k in tensor_list[0].keys(): | |
| out_dict[k] = [t[k] for t in tensor_list] | |
| out_dict[k] = pad_sequence( | |
| out_dict[k], | |
| padding_value=np.nan, | |
| batch_first=True, | |
| ) | |
| return out_dict | |