oneflow.nn.functional.affine_grid(theta, size: List[int], align_corners: bool = False)

The interface is consistent with PyTorch. The documentation is referenced from:

Generates a 2D or 3D flow field (sampling grid), given a batch of affine matrices theta.


This function is often used in conjunction with grid_sample() to build Spatial Transformer Networks .

  • theta (Tensor) – input batch of affine matrices with shape (\(N, 2, 3\)) for 2D or (\(N, 3, 4\)) for 3D

  • size (oneflow.Size) – the target output image size. (\(N, C, H, W\) for 2D or \(N, C, D, H, W\) for 3D) Example: oneflow.Size((32, 3, 24, 24))

  • align_corners (bool) – if True, consider -1 and 1 to refer to the centers of the corner pixels rather than the image corners. Refer to grid_sample() for a more complete description. A grid generated by affine_grid() should be passed to grid_sample() with the same setting for this option. Default: False


output Tensor of size (\(N, H, W, 2\))

Return type

output (Tensor)


>>> import oneflow as flow
>>> import numpy as np
>>> input = flow.tensor(np.arange(1., 7).reshape((1, 2, 3)), dtype=flow.float32)
>>> output = flow.nn.functional.affine_grid(input, flow.Size([1, 1, 2, 2]), align_corners=True)
>>> output
tensor([[[[ 0., -3.],
          [ 2.,  5.]],

         [[ 4.,  7.],
          [ 6., 15.]]]], dtype=oneflow.float32)