File size: 1,907 Bytes
6aab6b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | from typing import List, Tuple
import torch
from torch import Tensor
from torch.nn.functional import pad
def pad_dim(data: Tensor, dim: int, pad_len: float, value: float = 0) -> Tensor:
"""Pad a tensor along a given dimension.
Parameters
----------
data : Tensor
The input tensor.
dim : int
The dimension to pad.
pad_len : float
The padding length.
value : int, optional
The value to pad with.
Returns
-------
Tensor
The padded tensor.
"""
if pad_len == 0:
return data
total_dims = len(data.shape)
padding = [0] * (2 * (total_dims - dim))
padding[2 * (total_dims - 1 - dim) + 1] = pad_len
return pad(data, tuple(padding), value=value)
def pad_to_max(data: List[Tensor], value: float = 0) -> Tuple[Tensor, Tensor]:
"""Pad the data in all dimensions to the maximum found.
Parameters
----------
data : List[Tensor]
List of tensors to pad.
value : float
The value to use for padding.
Returns
-------
Tensor
The padded tensor.
Tensor
The padding mask.
"""
if isinstance(data[0], str):
return data, 0
# Check if all have the same shape
if all(d.shape == data[0].shape for d in data):
return torch.stack(data, dim=0), 0
# Get the maximum in each dimension
num_dims = len(data[0].shape)
max_dims = [max(d.shape[i] for d in data) for i in range(num_dims)]
# Get the padding lengths
pad_lengths = []
for d in data:
dims = []
for i in range(num_dims):
dims.append(0)
dims.append(max_dims[num_dims - i - 1] - d.shape[num_dims - i - 1])
pad_lengths.append(dims)
data = [pad(d, pad_len, value=value) for d, pad_len in zip(data, pad_lengths)]
# Stack the data
data = torch.stack(data, dim=0)
return data
|