conv_mc_BP.py

import os
import torch
import argparse
import datetime
import numpy as np
import torch.nn as nn
import torch.nn.functional as F

from utils import log
from random import randrange
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, SubsetRandomSampler


class ToeplitzConv2D(nn.Module):
    def __init__(self, in_channels, out_channels, input_size, filter_size):
        super(ToeplitzConv2D, self).__init__()
        # -- set params
        self.filter_h, self.filter_w = filter_size
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.input_size = input_size

        # -- initialize kernel
        self.kernel = torch.Tensor(out_channels, in_channels, self.filter_h, self.filter_w)
        nn.init.xavier_uniform_(self.kernel)

        # -- construct Toeplitz matrix
        self.weight = nn.Parameter(self.construct_toeplitz_2d_multi_channel())
        self.weight.mask = self.weight.ne(0).int()

    def construct_toeplitz_2d_multi_channel(self):
        """
            Construct a Toeplitz matrix for 2D convolution with multiple channels.
        """
        # -- get dimensions
        input_h, input_w = self.input_size
        output_h = input_h - self.filter_h + 1
        output_w = input_w - self.filter_w + 1
        self.out_dim = output_h * output_w

        # -- initialize Toeplitz matrix
        toeplitz_matrix = torch.zeros((self.out_channels, self.out_dim, self.in_channels * input_h * input_w))

        # -- construct Toeplitz matrix
        for out_ch in range(self.out_channels):
            for in_ch in range(self.in_channels):
                for i in range(output_h):
                    for j in range(output_w):
                        row = i * output_w + j
                        for m in range(self.filter_h):
                            for n in range(self.filter_w):
                                col = in_ch * input_h * input_w + (i + m) * input_w + (j + n)
                                toeplitz_matrix[out_ch, row, col] = self.kernel[out_ch, in_ch, m, n]

        return toeplitz_matrix

    def forward(self, x):
        batch_size = x.size(0)
        input_flat = x.view(batch_size, -1)

        # -- forward pass
        output = torch.stack([torch.matmul(input_flat, self.weight[j].T) for j in range(self.out_channels)], dim=1)
        # fixme: why not .view(B, C, output_h, output_w) here?
        # return output.view(batch_size, self.out_channels, int(self.out_dim ** 0.5), int(self.out_dim ** 0.5))
        return output


class MyModel(nn.Module):
    def __init__(self, args, in_channel, input_size, out_dim=10):
        super(MyModel, self).__init__()
        """
            Initialize MyModel object.
        """
        self.L = len(args.hidden_dims)

        # -- forward pathway
        for i, (out_channel, filter_size) in enumerate(args.conv_layers_config):
            setattr(self, f'tpz{i + 1}', ToeplitzConv2D(in_channels=in_channel, out_channels=out_channel,
                                                        input_size=input_size, filter_size=(filter_size, filter_size)))
            in_channel = out_channel
            input_size = (input_size[0] - filter_size + 1, input_size[1] - filter_size + 1)

        prev_dim = in_channel * input_size[0] * input_size[1]
        for i, hidden_dim in enumerate(args.hidden_dims):
            setattr(self, f'fc{i + 1}', nn.Linear(prev_dim, hidden_dim, bias=False))
            prev_dim = hidden_dim
        setattr(self, f'fc{len(args.hidden_dims) + 1}', nn.Linear(prev_dim, out_dim, bias=False))

        self.sopl = nn.Softplus(beta=10)

    def forward(self, x):

        y = [x]
        for name, layer in self.named_children():
            if 'tpz' in name:
                y.append(self.sopl(layer(y[-1])))
            elif 'fc' in name and str(1) in name:
                y.append(self.sopl(layer(y[-1].view(y[-1].size(0), -1))))
            elif 'fc' in name and str(self.L + 1) not in name:
                y.append(self.sopl(layer(y[-1])))
            elif 'fc' in name:
                logit = layer(y[-1])

        return y, logit


def stats(args, model, loss_func, test_loader):
    """
    Compute statistics
    """
    with torch.no_grad():
        x, test_label = next(iter(test_loader))

        x = x.view(x.size(0), -1)
        hidden, logits = model(x)

        # -- fetch activation
        activation = [*hidden, F.softmax(logits, dim=1)]

        # -- compute accuracy
        pred = F.softmax(logits, dim=1).argmax(dim=1)
        acc = torch.eq(pred, test_label).sum().item() / len(test_label)
        log([acc], f'{args.res_dir}/acc.txt')

        # -- compute loss
        loss = loss_func(logits, test_label)
        log([loss.item()], f'{args.res_dir}/loss.txt')

        # -- compute Oja's reconstruction error
        rec_err_a = []
        for idx, (name, m) in enumerate(model.named_children()):
            if isinstance(m, ToeplitzConv2D):
                for out_ch in range(m.out_channels):
                    rec_err_a.append(((activation[idx].view(activation[idx].size(0), -1) -
                                       torch.matmul(activation[idx + 1][:, out_ch, :],
                                                    m.weight[out_ch])).norm(dim=1) ** 2).mean().item())
            elif isinstance(m, nn.Linear):
                rec_err_a.append(((activation[idx].view(activation[idx].size(0), -1) -
                                   torch.matmul(activation[idx + 1], m.weight)).norm(dim=1) ** 2).mean().item())
        log(rec_err_a, f'{args.res_dir}/rec_err_a.txt')

        # -- compute activation statistics
        norm = [item for sublist in [y.norm(dim=y.dim() - 1).mean(dim=0).tolist() for y in activation] for item in
                (sublist if isinstance(sublist, list) else [sublist])]
        log(norm, f'{args.res_dir}/norm.txt')

        return acc, loss.item()


def parse_args():
    """
        Parses the input arguments
    """
    desc = "Pytorch implementation."
    parser = argparse.ArgumentParser(description=desc)

    # -- set model params
    parser.add_argument('--out_channels', type=int, nargs='+', default=[2, 3, 3],
                        help='Number of output channels.')
    parser.add_argument('--filter_size', type=int, nargs='+', default=[5, 5, 5],
                        help='Filter size for convolutional layers.')
    parser.add_argument('--hidden_dims', type=int, nargs='+', default=[170, 100, 70, 30],
                        help='Hidden dimensions for linear layers.')

    # -- set processor params
    parser.add_argument('--gpu_mode', type=int, default=1, help='Accelerate the script using GPU.')

    # -- set training params
    parser.add_argument('--seed', type=int, default=5, help='Random seed.')
    parser.add_argument('--epochs', type=int, default=1, help='Number of training epochs.')
    parser.add_argument('--n_train', type=int, default=5000, help='.')
    parser.add_argument('--n_test', type=int, default=100, help='.')
    parser.add_argument('--database', type=str, default='FashionMNIST', help='Training database.')
    parser.add_argument('--dim', type=float, nargs=2, default=[28, 28], help='Input dimension.')

    # -- set plasticity params
    parser.add_argument('--Theta', type=float, nargs='*', default=[0.001], help='Plasticity hyper-parameters.')

    # -- set save directory
    parser.add_argument('--test_name', type=str, default='Toeplitz_conv2D', help='.')
    parser.add_argument('--test_sub_name', type=str, default='Test_1', help='.')

    args = parser.parse_args()

    # -- set convolutional layers config
    args.conv_layers_config = list(zip(args.out_channels, args.filter_size))

    # -- GPU settings
    args.device = torch.device('cuda' if (bool(args.gpu_mode) and torch.cuda.is_available()) else 'cpu')

    # -- set results directory
    args.res_dir = f'./results/{args.test_name}/{args.test_sub_name}/' \
        f'{datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")}_{str(randrange(80))}'

    if not os.path.exists(args.res_dir):
        os.makedirs(args.res_dir)

    # -- store settings
    with open(f'{args.res_dir}/args.txt', 'w') as fp:
        for item in vars(args).items():
            fp.write(f'{item[0]} : {item[1]}\n')

    return args


def main():

    # -- parse input arguments
    args = parse_args()

    # -- set seeds
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    # -- load data
    if os.path.exists('../../data'):
        data_dir = '../../data'
    elif os.path.exists('../MetaPlasticity/data'):
        data_dir = '../MetaPlasticity/data'
    else:
        raise ValueError('data directory not found')

    transform = transforms.Compose([transforms.Resize((args.dim[0], args.dim[1])), transforms.ToTensor()])

    if args.database == 'CIFAR10':
        in_channel = 3
        dataset_train = datasets.CIFAR10(data_dir, train=True, download=False, transform=transform)
        dataset_test = datasets.CIFAR10(data_dir, train=False, download=False, transform=transform)
    elif args.database == 'FashionMNIST':
        in_channel = 1
        dataset_train = datasets.FashionMNIST(data_dir, train=True, download=False, transform=transform)
        dataset_test = datasets.FashionMNIST(data_dir, train=False, download=False, transform=transform)
    elif args.database == 'MNIST':
        in_channel = 1
        dataset_train = datasets.MNIST(data_dir, train=True, download=False, transform=transform)
        dataset_test = datasets.MNIST(data_dir, train=False, download=False, transform=transform)
    else:
        raise ValueError(f'Unknown database: {args.database}')

    train_sampler = SubsetRandomSampler(np.random.choice(range(50000), args.n_train, False))
    test_sampler = SubsetRandomSampler(np.random.choice(range(10000), args.n_test, False))
    train_loader = DataLoader(dataset_train, batch_size=1, sampler=train_sampler)
    test_loader = DataLoader(dataset_test, batch_size=args.n_test, sampler=test_sampler)

    # -- load model
    model = MyModel(args, in_channel=in_channel, input_size=args.dim)

    # -- optimizer settings
    loss_func = nn.CrossEntropyLoss()
    optim = torch.optim.Adam(model.parameters(), lr=args.Theta[0])

    # -- training loop
    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
        hidden, logits = model(input_tensor_flat)
        activation = [*hidden, F.softmax(logits, dim=1)]

        # -- compute loss
        loss = loss_func(logits, labels)

        # -- backpropagation
        optim.zero_grad()
        loss.backward()

        # -- mask toeplitz weight gradients
        for idx, (name, m) in enumerate(model.named_modules()):
            if isinstance(m, ToeplitzConv2D):
                for out_ch in range(m.out_channels):
                    m.weight.grad[out_ch] *= m.weight.mask[out_ch]

        # -- apply Oja's subspace rule
        if len(args.Theta) > 1:
            for idx, (name, m) in enumerate(model.named_children()):
                if isinstance(m, ToeplitzConv2D):
                    for out_ch in range(m.out_channels):
                        m.weight.data[out_ch] += \
                            (args.Theta[1] * (torch.matmul(activation[idx+1][:, out_ch, :].T,
                                                           activation[idx].view(activation[idx].size(0), -1)) -
                                              torch.matmul(torch.matmul(activation[idx+1][:, out_ch, :].T,
                                                                        activation[idx+1][:, out_ch, :]),
                                                           m.weight.data[out_ch])))
                        m.weight.data[out_ch] *= m.weight.mask[out_ch]
                elif isinstance(m, nn.Linear):
                    m.weight.data += args.Theta[1] * (torch.matmul(activation[idx+1].T, activation[idx].view(activation[idx].size(0), -1)) -
                                                      torch.matmul(torch.matmul(activation[idx+1].T, activation[idx+1]),
                                                                   m.weight.data))

        # -- update weights
        optim.step()

        # -- log statistics
        acc, loss = stats(args, model, loss_func, test_loader)
        print(f'Epoch {i + 1}, Loss: {loss:.4f}, acc: {acc}')


if __name__ == '__main__':
    main()