Downsample a stack of 2d images in PyTorch

downsample.py

def downsample_2d(X, sz):
    """
    Downsamples a stack of square images.
    
    Args:
        X: a stack of images (batch, channels, ny, ny).
        sz: the desired size of images.
        
    Returns:
        The downsampled images, a tensor of shape (batch, channel, sz, sz)
    """
    kernel = torch.tensor([[.25, .5, .25], 
                           [.5, 1, .5], 
                           [.25, .5, .25]], device=X.device).reshape(1, 1, 3, 3)
    kernel = kernel.repeat((X.shape[1], 1, 1, 1))
    while sz < X.shape[-1] / 2:
        # Downsample by a factor 2 with smoothing
        mask = torch.ones(1, *X.shape[1:])
        mask = F.conv2d(mask, kernel, groups=X.shape[1], stride=2, padding=1)
        X = F.conv2d(X, kernel, groups=X.shape[1], stride=2, padding=1)
        
        # Normalize the edges and corners.
        X = X = X / mask
    
    return F.interpolate(X, size=sz, mode='bilinear')