File size: 592 Bytes
bd3493c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | 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
|