def _norm_except_dim(v: Tensor, dim: int):
assert -v.dim() <= dim <= v.dim() - 1, "dim out of range"
if dim == -1:
return flow.linalg.norm(v, ord="fro")
elif dim == 0:
return _norm_except_dim_0(v)
elif dim == v.dim() - 1:
output_size = [1] * v.dim()
output_size[v.dim() - 1] = v.size(v.dim() - 1)
return flow.linalg.norm(v.view(-1, v.size(v.dim() - 1)), ord=2,
dim=0).view(*output_size)
else:
return flow.transpose(_norm_except_dim_0(flow.transpose(v, 0, dim)), 0,
dim)
def forward(self, input: Tensor) -> Tensor:
assert len(input.shape) == 4, "Only Accept 4D Tensor"
(_batch, _channel, _height, _width) = input.shape
assert (
_channel % (self.h_upscale_factor * self.w_upscale_factor) == 0
), "The channels of input tensor must be divisible by (upscale_factor * upscale_factor) or (h_upscale_factor * w_upscale_factor)"
_new_c = int(_channel / (self.h_upscale_factor * self.w_upscale_factor))
out = input.reshape(
_batch,
_new_c,
self.h_upscale_factor * self.w_upscale_factor,
_height,
_width,
)
out = out.reshape(
_batch,
_new_c,
self.h_upscale_factor,
self.w_upscale_factor,
_height,
_width,
)
out = out.permute(0, 1, 4, 2, 5, 3)
out = out.reshape(
_batch,
_new_c,
_height * self.h_upscale_factor,
_width * self.w_upscale_factor,
)
return out
def unpad_sequence(
padded_sequences: Tensor,
lengths: Tensor,
batch_first: bool = False,
) -> List[Tensor]:
"""
Unpad padded Tensor into a list of variable length Tensors
``unpad_sequence`` unstacks padded Tensor into a list of variable length Tensors.
Args:
padded_sequences (Tensor): padded sequences.
lengths (Tensor): length of original (unpadded) sequences.
batch_first (bool, optional): whether batch dimension first or not. Default: False.
Returns:
a list of :class:`Tensor` objects
For example:
.. code-block:: python
>>> from oneflow.nn.utils.rnn import pad_sequence, unpad_sequence
>>> import oneflow as flow
>>> import numpy as np
>>> a = flow.ones(25, 300)
>>> b = flow.ones(22, 300)
>>> c = flow.ones(15, 300)
>>> sequences = [a, b, c]
>>> padded_sequences = pad_sequence(sequences)
>>> lengths = flow.as_tensor([v.size(0) for v in sequences])
>>> unpadded_sequences = unpad_sequence(padded_sequences, lengths)
>>> np.allclose(sequences[0].numpy(), unpadded_sequences[0].numpy())
True
>>> np.allclose(sequences[1].numpy(), unpadded_sequences[1].numpy())
True
>>> np.allclose(sequences[2].numpy(), unpadded_sequences[2].numpy())
True
"""
unpadded_sequences = []
if not batch_first:
padded_sequences = padded_sequences.permute((1, 0, 2))
max_length = padded_sequences.shape[1]
idx = flow.arange(max_length)
for seq, length in zip(padded_sequences, lengths):
mask = idx < length
unpacked_seq = seq[mask]
unpadded_sequences.append(unpacked_seq)
return unpadded_sequences
def pack_padded_sequence(
input: Tensor,
lengths: Tensor,
batch_first: bool = False,
enforce_sorted: bool = True,
) -> PackedSequence:
"""The interface is consistent with PyTorch.
The documentation is referenced from: https://pytorch.org/docs/stable/generated/torch.nn.utils.rnn.pack_padded_sequence.html.
Packs a Tensor containing padded sequences of variable length.
:attr:`input` can be of size ``T x B x *`` where `T` is the length of the
longest sequence (equal to ``lengths[0]``), ``B`` is the batch size, and
``*`` is any number of dimensions (including 0). If ``batch_first`` is
``True``, ``B x T x *`` :attr:`input` is expected.
For unsorted sequences, use `enforce_sorted = False`. If :attr:`enforce_sorted` is
``True``, the sequences should be sorted by length in a decreasing order, i.e.
``input[:,0]`` should be the longest sequence, and ``input[:,B-1]`` the shortest
one. `enforce_sorted = True` is only necessary for ONNX export.
Note:
This function accepts any input that has at least two dimensions. You
can apply it to pack the labels, and use the output of the RNN with
them to compute the loss directly. A Tensor can be retrieved from
a :class:`PackedSequence` object by accessing its ``.data`` attribute.
Args:
input (Tensor): padded batch of variable length sequences.
lengths (Tensor or list(int)): list of sequence lengths of each batch
element (must be on the CPU if provided as a tensor).
batch_first (bool, optional): if ``True``, the input is expected in ``B x T x *``
format.
enforce_sorted (bool, optional): if ``True``, the input is expected to
contain sequences sorted by length in a decreasing order. If
``False``, the input will get sorted unconditionally. Default: ``True``.
Returns:
a :class:`PackedSequence` object
"""
lengths = flow.as_tensor(lengths, dtype=flow.int64)
assert (
enforce_sorted == True
), "Only support enforce_sorted == True for now. Plesase Sort the input by length in a decreasing order."
if enforce_sorted:
sorted_indices = None
else:
lengths, sorted_indices = flow.sort(lengths, descending=True)
sorted_indices = sorted_indices.to(input.device)
batch_dim = 0 if batch_first else 1
input = input.index_select(batch_dim, sorted_indices)
data, batch_sizes = flow._C.pack_padded_sequence(input, lengths,
batch_first)
return PackedSequence(data, batch_sizes, sorted_indices, None)
def grad(
outputs: Union[Tensor, Sequence[Tensor]],
inputs: Union[Tensor, Sequence[Tensor]],
out_grads: Union[Tensor, Sequence[Tensor], None] = None,
retain_graph: bool = False,
create_graph: bool = False,
) -> Tuple[Tensor]:
in_grads = grad_api(
convert_to_tensor_tuple(outputs),
convert_to_tensor_tuple(inputs),
convert_to_tensor_tuple(out_grads),
retain_graph,
create_graph,
)
return tuple([Tensor(x) for x in in_grads])
def grad(
outputs: Union[Tensor, Sequence[Tensor]],
inputs: Union[Tensor, Sequence[Tensor]],
grad_outputs: Union[Tensor, Sequence[Tensor], None] = None,
retain_graph: bool = False,
create_graph: bool = False,
) -> Tuple[Tensor]:
r"""
The documentation is referenced from:
https://pytorch.org/docs/1.10/generated/torch.autograd.grad.html.
Computes and returns the sum of gradients of outputs with respect to the inputs.
The graph is differentiated using the chain rule. ``grad_outputs`` should be a sequence of
length matching ``outputs``, containing the "vector" in the Jacobian-vector product.
(``None`` is an acceptable value for that tensor don't require gradient.)
Args:
outputs (Sequence[Tensor]): Tensors of which the derivative will be computed.
inputs (Sequence[Tensor]): Inputs w.r.t. which the derivative will be returned(and not
accumulated into ``.grad``).
grad_outputs (Sequence[Tensor], optional): The "vector" in the Jacobian-vector product.
Usually gradients w.r.t. each output. None values can be specified for scalar Tensors
or ones that don't require grad. Defaults to None.
retain_graph (bool, optional): If ``False``, the graph used to compute the grads will be
reset after backward is complete. Defaults to ``False``. Note that in nearly all cases
setting this option to ``True`` is not needed and often can be worked around in a much
more efficient way. Defaults to the value of ``create_graph``.
create_graph (bool, optional): If ``True``, graph of the derivative will be constructed,
allowing to compute higher order derivative products. Defaults to ``False``.
Returns:
Tuple(Tensor): A tuple of tensors containing the gradients for each ``inputs``.
"""
in_grads = grad_api(
convert_to_tensor_tuple(outputs),
convert_to_tensor_tuple(inputs),
convert_to_tensor_tuple(grad_outputs),
retain_graph,
create_graph,
)
return tuple([Tensor(x) for x in in_grads])
def forward(
self,
log_probs: Tensor,
targets: Tensor,
input_lengths: Tensor,
target_lengths: Tensor,
) -> Tensor:
max_target_length = 0
if targets.ndim == 1:
max_target_length = target_lengths.max().item()
elif targets.ndim == 2:
max_target_length = targets.shape[1]
return flow._C.ctc_loss(
log_probs,
targets,
input_lengths,
target_lengths,
max_target_length,
self.blank,
self.zero_infinity,
self.reduction,
)
def __init__(
self,
num_embeddings: int,
embedding_dim: int,
padding_idx: Optional[int] = None,
max_norm: Optional[float] = None,
norm_type: float = 2.0,
scale_grad_by_freq: bool = False,
sparse: bool = False,
_weight: Optional[Tensor] = None,
):
super().__init__()
self.num_embeddings = num_embeddings
self.embedding_dim = embedding_dim
if padding_idx is not None:
if padding_idx > 0:
assert (padding_idx < self.num_embeddings
), "Padding_idx must be within num_embeddings"
elif padding_idx < 0:
assert (padding_idx >= -self.num_embeddings
), "Padding_idx must be within num_embeddings"
padding_idx = self.num_embeddings + padding_idx
self.padding_idx = padding_idx
self.max_norm = max_norm
self.norm_type = norm_type
self.scale_grad_by_freq = scale_grad_by_freq
assert sparse is False, "Not support sparse=True yet!"
if _weight is None:
self.weight = flow.nn.Parameter(
Tensor(num_embeddings, embedding_dim))
self.reset_parameters()
else:
assert list(_weight.shape) == [
num_embeddings,
embedding_dim,
], "Shape of weight does not match num_embeddings and embedding_dim"
self.weight = flow.nn.Parameter(_weight)
self.sparse = sparse
def _norm_except_dim_0(v: Tensor):
output_size = [1] * v.dim()
output_size[0] = v.size(0)
return flow.linalg.norm(v.view(v.size(0), -1), ord=2,
dim=1).view(*output_size)