conv_mc.py

import torch
import torch.nn as nn
import numpy as np
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, SubsetRandomSampler
import matplotlib.pyplot as plt


def construct_toeplitz_2d_multi_channel(input_size, filter_weights):
    """
    Construct a Toeplitz matrix for 2D convolution with multiple channels.
    """
    input_h, input_w = input_size
    out_channels, in_channels, filter_h, filter_w = filter_weights.shape
    output_h = input_h - filter_h + 1
    output_w = input_w - filter_w + 1

    toeplitz_matrix = torch.zeros((out_channels, output_h * output_w, in_channels * input_h * input_w))

    for out_ch in range(out_channels):
        for in_ch in range(in_channels):
            for i in range(output_h):
                for j in range(output_w):
                    row = i * output_w + j
                    for m in range(filter_h):
                        for n in range(filter_w):
                            col = in_ch * input_h * input_w + (i + m) * input_w + (j + n)
                            toeplitz_matrix[out_ch, row, col] = filter_weights[out_ch, in_ch, m, n]

    return toeplitz_matrix


def model(input, w):
    """
    Forward pass
    """
    # -- set params
    sopl = nn.Softplus(beta=10)
    batch_size = input.size(0)
    out_channels, out_dim, _ = w.size()

    # -- forward pass
    output_flat = torch.zeros(batch_size, out_channels, out_dim)
    for out_ch in range(out_channels):
        output_flat[:, out_ch, :] = torch.matmul(input, w[out_ch].T)

    return sopl(output_flat)


def stats(w, test_loader):
    """
    Compute statistics
    """
    # -- load test data
    x, test_label = next(iter(test_loader))

    # -- resize input
    out_channels, out_dim, _ = w.size()
    batch_size, in_channels, input_height, input_width = x.size()
    x = x.view(batch_size, in_channels * input_height * input_width)

    # -- forward pass
    y = model(x, w)

    # -- compute reconstruction error
    rec_err_a = []
    for out_ch in range(out_channels):
        rec_err_a.append(((x - torch.matmul(y[:, out_ch, :], w[out_ch])).norm(dim=1) ** 2).mean().item())
        # rec_err_b = ((torch.matmul(x, w.permute(0, 2, 1)) - torch.matmul(y,
        # torch.matmul(w, w.permute(0, 2, 1)))).norm(dim=1) ** 2).mean().item()

    # -- compute activation statistics
    y_norm_flat = [y_.norm(dim=1).mean().item() for y_ in [x, y.view(batch_size, -1)]]
    y_norm_chnl = [y_.norm(dim=1).mean().item() for y_ in [x, *[y[:, out_ch, :] for out_ch in range(out_channels)]]]

    return rec_err_a, y_norm_flat, y_norm_chnl

# -- params
Theta = 0.01
n_train, n_test = 5000, 10

# -- load data
inp_h, inp_w = 7, 7
data_dir = '../../data'
transform = transforms.Compose([transforms.Resize((inp_h, inp_w)), transforms.ToTensor()])
dataset_train = datasets.MNIST(data_dir, train=True, download=False, transform=transform)
dataset_test = datasets.MNIST(data_dir, train=False, download=False, transform=transform)
train_sampler = SubsetRandomSampler(np.random.choice(range(50000), n_train, False))
test_sampler = SubsetRandomSampler(np.random.choice(range(10000), n_test, False))
train_loader = DataLoader(dataset_train, batch_size=1, sampler=train_sampler)
test_loader = DataLoader(dataset_test, batch_size=n_test, sampler=test_sampler)

# -- construct filter
in_channels, out_channels, filter_h, filter_w = 1, 2, 3, 3
conv_layer = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(filter_h, filter_w), bias=False)
toeplitz_matrix = construct_toeplitz_2d_multi_channel((inp_h, inp_w), conv_layer.weight.data.clone())
mask = toeplitz_matrix.ne(0).int()

# -- training loop
rec_err_a_ch1, rec_err_a_ch2, y_norm_flat, y_norm_chnl = [], [], [], []
for i, (input_tensor, labels) in enumerate(train_loader):
    # -- resize input
    batch_size, in_channels, input_height, input_width = input_tensor.size()
    input_tensor_flat = input_tensor.view(batch_size, in_channels * input_height * input_width)

    # -- forward pass
    output = model(input_tensor_flat, toeplitz_matrix)

    # -- apply Oja's subspace rule
    for out_ch in range(out_channels):
        toeplitz_matrix[out_ch] += Theta * (torch.matmul(output[:, out_ch, :].T, input_tensor_flat) -
                                            torch.matmul(torch.matmul(output[:, out_ch, :].T, output[:, out_ch, :]),
                                                         toeplitz_matrix[out_ch]))
        toeplitz_matrix[out_ch] *= mask[out_ch]

    # -- compute statistics
    my_stats = stats(toeplitz_matrix, test_loader)
    rec_err_a_ch1.append(my_stats[0][0])
    rec_err_a_ch2.append(my_stats[0][1])
    y_norm_flat.append(my_stats[1])
    y_norm_chnl.append(my_stats[2])

plt.plot(rec_err_a_ch1, label='Reconstruction Error Channel 1')
plt.show()
plt.close()
plt.plot(rec_err_a_ch2, label='Reconstruction Error Channel 2')
plt.show()
plt.close()
plt.plot(np.array(y_norm_flat)[:, 0], label='Activation Norm Flat')
plt.plot(np.array(y_norm_flat)[:, 1], label='Activation Norm Flat')
plt.show()
plt.close()
plt.plot(np.array(y_norm_chnl)[:, 0], label='Activation Norm per Channel')
plt.plot(np.array(y_norm_chnl)[:, 1], label='Activation Norm per Channel')
plt.plot(np.array(y_norm_chnl)[:, 2], label='Activation Norm per Channel')
plt.show()
plt.close()