PyTorch code snippets

Date written Jul 5, 2020
Date updated Last updated: Dec 31, 2020
Filed under Software Eng. in tech


CUDA Specific Builds

To install specific CUDA version builds for PyTorch, use

$ pip install torch -f

cu101 corresponds to CUDA 10.1, and needs to be changed appropriately.

Vectorized Pairwise Distances

For XR...×m×d,YR...×n×d\mathbf{X} \in \mathbb{R}^{... \times m \times d}, \mathbf{Y} \in \mathbb{R}^{... \times n \times d}, the pairwise distance matrix between each pair of these batched matrices is DR...×m×n\mathbf{D} \in \mathbb{R}^{... \times m \times n}, where ... represent arbitrary batch dimension (think batches of pairs of mm and nn samples of dimension dd).

def pairwise_dist(x, y):
xx = (x * x).sum(dim=-1).unsqueeze(-1)
yy = (y * y).sum(dim=-1).unsqueeze(-2)
xy = torch.einsum('...ji,>...jk', x, y)
d = xx + yy - 2. * xy
return d