lite_boost.layers.nearest_exact_upsample

View Source On AtomGit
lite_boost.layers.nearest_exact_upsample(x, size=None, scale_factor=None)[source]

Upsamples x using nearest-exact interpolation.

Behaves like torch.nn.functional.interpolate(mode="nearest-exact"). All supported dtypes run natively on A2 and other SoCs; on 300I Duo, bfloat16 inputs are computed through a float32 intermediate cast, which produces bitwise-identical results as nearest-exact uses integer indexing.

Parameters:

  • x (Tensor) – Input tensor with shape \((B, C, H, W)\). Supported dtypes are float16, float32 and bfloat16.

  • size (Union[int, tuple[int]], optional) – Output spatial size. Provide exactly one of size and scale_factor. Default: None.

  • scale_factor (Union[float, tuple[float]], optional) – Multiplier for the spatial dims. Provide exactly one of size and scale_factor. Default: None.

Returns:

Tensor, with the same dtype as x, and shape determined by size or scale_factor.

Raises:

  • ValueError – If x is not 4-D.

  • ValueError – If both or neither of size and scale_factor are provided.

Supported Platforms:

Ascend

Examples

>>> import torch
>>> import torch_npu
>>> from lite_boost.layers import nearest_exact_upsample
>>> torch.npu.set_device(0)
>>> x = torch.arange(4, device="npu").view(1, 1, 2, 2).float()
>>> y = nearest_exact_upsample(x, scale_factor=2)
>>> print(y.shape)
torch.Size([1, 1, 4, 4])
>>> print(y[0, 0, 0])
tensor([0., 0., 1., 1.], device='npu:0')