affine_grid() in pytorch

According to PyTorch document, In PyTorch, torch.nn.functional.affine_grid generates a 2D or 3D flow field (sampling grid) based on a batch of affine matrices. It is almost exclusively used in combination with torch.nn.functional.grid_sample to perform geometric transformations like rotation, translation, scaling, and shearing.

The following code examples show part of a Python class, which implements augmentation to 2D data, in which affine_grid() is applied.

class SegmentationAugmentation(nn.Module):
    def __init__(
            self, flip=None, offset=None, scale=None, rotate=None, noise=None
    ):
        super().__init__()

        self.flip = flip
        self.offset = offset
        self.scale = scale
        self.rotate = rotate
        self.noise = noise

    def forward(self, input_g, label_g):
        #transform_t is 3*3 matrix
        transform_t = self._build2dTransformMatrix()
        transform_t = transform_t.expand(input_g.shape[0], -1, -1)
        transform_t = transform_t.to(input_g.device, torch.float32)
        #The first dimension of the transformation is the batch,but 
        #we only want the first two rows of the 3  × 3 matrices per
        #batch item.
        affine_t = F.affine_grid(transform_t[:,:2],
                input_g.size(), align_corners=False)

        augmented_input_g = F.grid_sample(input_g,
                affine_t, padding_mode='border',
                align_corners=False)
        #We need the same transformation applied to data and
        #label, so we use the same grid.
        augmented_label_g = F.grid_sample(label_g.to(torch.float32),
                affine_t, padding_mode='border',
                align_corners=False)