xonv: Extended convolutional layers
Project description
xonv: Extended convolutional layers
This repository contains the code for extended convolutional layers. These layers are akin to the convolutional layers in PyTorch, but with the key difference that they have spatially varying kernels.
Since the kernels are spatially varying, the convolutional layers in this repository offer more expressive power than the standard convolutional layers while having $\mathcal{O}(n)$ parameters, where $n$ is the input size. The implementation is based on matrix-vector products, which allows for scalable training and inference on GPUs.
Below is a comparison between the toeplitz-like matrix associated with a
regular convolutional layer vs the extended convolutional layer
(xonv
):
Installation
Run the command below to install the package to be used in your Python environment.
pip install xonv
For further development and to run the examples, clone the repository
and install the package in editable mode. Make sure to adapt CUDA
version in setup.cfg
to the one installed on your system.
# Create a new conda environment.
conda create --name xonv "python<=3.12"
conda activate xonv
# Clone the repository and install the package in editable mode.
git clone ttps://github.com/alisiahkoohi/xonv
cd xonv/
pip install -e .
Usage
The extended convolutional layers can be used as a drop-in replacement for the PyTorch convolutional layers. The following example demonstrates how to use the extended convolutional layers:
from xonv.layer import Xonv2D
input_size = (32, 32) # Height, Width of input
in_channels = 3
out_channels = 16
kernel_size = 3
stride = 2
layer = Xonv2D(
in_channels,
out_channels,
kernel_size,
input_size,
stride=stride,
)
input_tensor = torch.randn(1, in_channels, *input_size)
output = layer(input_tensor)
print(layer) # Xonv2D(in_channels=3, out_channels=16, kernel_size=3, input_size=(32, 32), stride=2)
print(output.shape) # Should be [1, 16, 16, 16]
Examples
Visualizing the associated linear matrix
To visualize the toeplitz-like matrix associated with the convolutional layer, run the following command:
python scripts/create_toeplitz_like_matrix.py
Comparing the loss landscape of Xonv2d
vs torch.nn.Conv2d
for a regression task
To compare the loss landscape of the extended convolutional layer with the standard convolutional layer, run the following command:
python scripts/regression_loss_landscape_comparison.py
Comparing convergence rate of Xonv2d
vs torch.nn.Conv2d
for a regression task
To compare the convergence rate of the extended convolutional layer with the standard convolutional layer, run the following command:
python scripts/regression_convergence_comparison.py
Questions
Please contact alisk@rice.edu for questions.
Author
Ali Siahkoohi
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.