Batch Affine Transform with PyTorch

You want to apply a single 2×3 affine transform to a batch of 2D points using PyTorch tensor ops only (no Python loops), by using homogeneous coordinates and a single matrix multiply.

---

1. Background Knowledge

An affine transformation in 2D is a function that maps a point $(x, y)$ to a new point $(x', y')$ by combining a linear transform (rotation, scaling, shear) and a translation:

$$\begin{pmatrix} x' \\ y' \end{pmatrix} = \begin{pmatrix} a & b & t_x \\ c & d & t_y \end{pmatrix} \begin{pmatrix} x \\ y \\ 1 \end{pmatrix}$$

The $2\times 2$ block $\begin{pmatrix} a & b \\ c & d \end{pmatrix}$ handles the linear part; $(t_x, t_y)$ is the translation.

To express this as a single matrix multiplication, we use homogeneous coordinates: we append a 1 to each 2D point so that each point becomes $(x, y, 1)$. For a batch of $N$ points, stacking them gives a tensor $P_h$ of shape $(N, 3)$. The affine matrix $M$ has shape $(2, 3)$. Using PyTorch’s row-major convention for batched points, the natural operation is:

$$P' = P_h \cdot M^T$$

where $P_h \in \mathbb{R}^{N\times3}$, $M^T \in \mathbb{R}^{3\times 2}$, giving $P' \in \mathbb{R}^{N\times 2}$.

In deep learning / CV code, this pattern shows up everywhere: you keep all points in a tensor of shape $(N, 2)$, convert once to homogeneous coordinates, and then use vectorized matrix multiplication so the operation is fast and GPU-friendly.

---

2. Algorithm / Approach

General pattern:

1. Input shapes:

• Points: tensor points of shape (N, 2) (each row [x, y]).
• Affine matrix: tensor M of shape (2, 3) (as in the formula).

2. Convert to homogeneous coordinates:

• Build a tensor ones of shape (N, 1) filled with 1.
• Concatenate: P_h = torch.cat([points, ones], dim=1) → shape (N, 3).

3. Matrix multiply:

• Use M.T (shape (3, 2)).
• P_prime = P_h @ M.T → shape (N, 2).

This is a standard “augment-then-matmul” pattern for applying affine transforms to batches of vectors.

---

3. Step-by-Step Strategy

1. Understand input tensors:

• Confirm points has shape (N, 2) and dtype (e.g., torch.float32).
• Confirm M has shape (2, 3) and same dtype/device.

2. Create homogeneous coordinates:

N = points.shape
ones = torch.ones(N, 1, device=points.device, dtype=points.dtype)
P_h = torch.cat([points, ones], dim=1) # (N, 3)

3. Apply affine transform:

M_T = M.transpose(0, 1) # (3, 2)
P_prime = P_h @ M_T # (N, 2)

4. Return result:

• Output P_prime as the transformed points.

5. (Optional) Batch of transforms:

• If you later extend to a batch of matrices (B, 2, 3) and points (B, N, 2), you will use broadcasting or einsum / bmm, but the idea remains: augment with ones, then multiply.

---

4. Common Pitfalls

• Shape mismatches:

• Using M instead of M.T: (N, 3) @ (2, 3) is invalid; you need (N, 3) @ (3, 2).

• Concatenating along wrong dimension (e.g., dim=0 instead of dim=1).

• Forgetting homogeneous coordinate:

• If you multiply (N, 2) by a (2, 2) linear matrix, you only get rotation/scale/shear, no translation.

• Device/dtype mismatch:

• Creating ones on CPU while points is on GPU, or using different dtypes; always match device and dtype.

• Incorrect point layout:

• Some APIs store as (2, N) instead of (N, 2). Ensure you know whether points are stored as rows or columns and adjust matmul accordingly.

---

5. Time & Space Complexity

Let $N$ be the number of points.

• Time complexity:

• Creating homogeneous coordinates: $O(N)$.

• Matrix multiplication $(N \times 3) \cdot (3 \times 2)$: $O(N)$.

• Overall: $O(N)$.

• Space complexity:

• Storing input points: $O(N)$.

• Homogeneous points P_h: $O(N)$.

• Output P_prime: $O(N)$.

• Overall additional memory beyond inputs: $O(N)$.