transposed_conv_torch.ipynb

Kernel: Python [conda env:pymc_exp]

Please find jax implementation of this notebook here: https://colab.research.google.com/github/probml/pyprobml/blob/master/notebooks/book1/14/transposed_conv_jax.ipynb

# Based on http://d2l.ai/chapter_computer-vision/transposed-conv.html
try:
    import torch
except ModuleNotFoundError:
    %pip install -qq torch
    import torch

from torch import nn
import numpy as np

def trans_conv(X, K):
    h, w = K.shape
    Y = torch.zeros((X.shape[0] + h - 1, X.shape[1] + w - 1))
    for i in range(X.shape[0]):
        for j in range(X.shape[1]):
            Y[i : i + h, j : j + w] += X[i, j] * K
    return Y


# Example from D2L fig 13.10.1
X = torch.tensor([[0.0, 1], [2, 3]])
K = torch.tensor([[0.0, 1], [2, 3]])
Y = trans_conv(X, K)
print(Y)

X, K = X.reshape(1, 1, 2, 2), K.reshape(1, 1, 2, 2)
tconv = nn.ConvTranspose2d(1, 1, kernel_size=2, bias=False)
tconv.weight.data = K
Y2 = tconv(X)
# print(Y2)
assert torch.allclose(Y, Y2)


"""
X, K = X.reshape(1, 1, 2, 2), K.reshape(1, 1, 2, 2)
tconv = nn.ConvTranspose2d(1, 1, kernel_size=2, padding = 1, bias=False)
tconv.weight.data = K
Y2 = tconv(X)
print('Y2', Y2)
"""


# Transposed Matrix multiplication

K = torch.tensor([[1, 2], [3, 4]])


def kernel2matrix(K):
    k, W = torch.zeros(5), torch.zeros((4, 9))
    k[:2], k[3:5] = K[0, :], K[1, :]
    W[0, :5], W[1, 1:6], W[2, 3:8], W[3, 4:] = k, k, k, k
    return W


W = kernel2matrix(K)

X = torch.tensor([[0.0, 1], [2, 3]])
Y = trans_conv(X, K)
Y2 = torch.mv(W.T, X.reshape(-1)).reshape(3, 3)
assert torch.allclose(Y, Y2)


# Example from Geron fig 14.27

X = torch.ones((2, 3))
K = torch.ones(3, 3)
X, K = X.reshape(1, 1, 2, 3), K.reshape(1, 1, 3, 3)
tconv = nn.ConvTranspose2d(1, 1, kernel_size=3, stride=2, bias=False)
tconv.weight.data = K
Y2 = tconv(X)
print(Y2.shape)

tensor([[ 0.,  0.,  1.],
        [ 0.,  4.,  6.],
        [ 4., 12.,  9.]])
torch.Size([1, 1, 5, 7])