YodaEmbedding
commented
10 months ago SpectralConv2d vs Conv2d mini-experiments
Below are a couple of example runs (with different randomly initialized kernels) to compare SpectralConv2d vs Conv2d. (Ignore the inaccurate titles; I accidentally set the title for everything to "smoothing kernel".)
sinusoidal channel averaging:
conv_kwargs = dict(in_channels=3, out_channels=2, kernel_size=11, padding=5)
def init_conv_target(conv):
k = 0.05 * torch.linspace(0, 2 * torch.pi, conv.kernel_size[0], device=device).sin()
conv.weight.data[:] = kSimple depthwise cosine "smoothing" kernel:
conv_kwargs = dict(in_channels=8, out_channels=8, kernel_size=11, padding=5)
def init_conv_target(conv):
k = torch.linspace(0, torch.pi, conv.kernel_size[0], device=device).sin()
k = k * k[:, None]
k = k / k.sum()
idx = torch.arange(conv.in_channels, device=device)
conv.weight.data[:] = 0
conv.weight.data[idx, idx, :, :] = kSimple depthwise vertical edge-detector kernel:
conv_kwargs = dict(in_channels=4, out_channels=4, kernel_size=5, padding=5, groups=4)
def init_conv_target(conv):
K = conv.kernel_size[0]
k = 2 * torch.linspace(-1, 1, K, device=device) / K**2
conv.weight.data[:] = kRandom initialization:
conv_kwargs = dict(in_channels=4, out_channels=4, kernel_size=5, padding=5)
def init_conv_target(conv):
pass # Maintain random initialization.Surprisingly, SpectralConv2d is not worse...!
Figures generated via:
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
from compressai.layers import SpectralConv2d
device = "cuda"
def train(conv, conv_target, max_steps=1000, lr=1e-4, batch_size=16):
losses = []
for step in range(max_steps):
conv.zero_grad()
x = torch.rand((batch_size, conv.in_channels, 256, 256)).to(device)
y = conv_target(x).detach()
y_hat = conv(x)
loss = ((y - y_hat) ** 2).mean()
loss.backward()
for param in conv.parameters():
param.data -= param.grad * lr
print(f"step {step:04d} loss {loss.item():.4f}")
losses.append(loss.item())
return losses
models = {
"spectral": SpectralConv2d(**conv_kwargs),
"regular": nn.Conv2d(**conv_kwargs),
}
# Initialize all models to exact same random initialization.
conv_rand = nn.Conv2d(**conv_kwargs).to(device)
init_weight = conv_rand.weight.data
init_bias = conv_rand.bias.data
for key, conv in models.items():
conv.to(device)
if isinstance(conv, SpectralConv2d):
conv.weight_transformed.data = conv._to_transform_domain(init_weight).clone()
else:
conv.weight.data = init_weight.clone()
conv.bias.data = init_bias.clone()
# Initialize "ideal"/target model kernels using the functions defined above.
conv_target = nn.Conv2d(**conv_kwargs).to(device)
init_conv_target(conv_target)
results = {key: train(conv, conv_target) for key, conv in models.items()}
fig, ax = plt.subplots()
for key, y in results.items():
ax.plot(y, label=key)
ax.legend()
ax.set(xlabel="step", ylabel="loss", title="Fitting to a target kernel")
fig.savefig("fit_spectralconv2d.png")Random musings:
I wonder why no one's applied this to other image problems (classification/superresolution/etc)? Those problems should also be concerned with "frequency regularized" kernels (so to speak).
I wonder if this could be turned into a short "paper" with more interesting experiments... Or if not, perhaps at least a "note" (1, 2) on arxiv.
Also, another related thought that I had once upon a time:
Generate the kernel from a different basis, e.g. the first $K=4$ 2D DCT basis elements. Assuming a single input and output channel, `y = conv(x, weight=k(w))`, where $w = (w_1, \ldots, w_K)$ are trainable weights and $k(w)$ generates the kernel via the weighted summation $k(w) = \sum_i w_i e_i$. ...Then, even large 7x7 or 9x9 kernels can be expressed by only a few parameters. Extend appropriately to multiple input/output channels. I guess this looks very related to "dynamic convolution" / "CondConv", but with 10x fewer parameters instead of 4x more parameters.
I think I like Balle's insight of using a big DFT and calling it a day, though. I guess one more thing we could do as an extension to Balle's "Spectral" reparameterization is something like:
```python
def _from_transform_domain(self, w: Tensor) -> Tensor:
# Attenuate out high-frequencies.
# Not sure if correctly written, but this is intended to keep the lower frequencies, and push the others to 0.
# d/dw is then [hopefully larger for the lower frequencies, but I should check if I wrote this correctly...].
yy = torch.linspace(1, 0.1, w.shape[-2], device=w.device)[None, :]
xx = torch.linspace(1, 0.1, w.shape[-1], device=w.device)
mask = yy * xx
w = w * mask
# Balle's original spectral reparametrization.
return torch.fft.irfftn(w, s=self.kernel_size, dim=self.dim, norm="ortho")
```
Or if high frequencies are also important, do this "regularizing" reparameterization for only some portion of the weights. That way, the network is forced into a balance of both, inhabits a ~~lower dimensional space~~ well-conditioned/"regularized" space (c.f. dynamic conv's affine constraint), and is roughly just as expressible as an unparameterized network.
TODO: Create a test suite of target kernels like the above, or pulled from trained networks, and see how quickly different `_from_transform_domain` functions converge... Is there something that fits our networks better than Balle's unmodified spectral conv?
fracape
Introduced in "Efficient Nonlinear Transforms for Lossy Image Compression" by Johannes Ballé, PCS 2018. Reparameterizes the weights to be derived from weights stored in the frequency domain. In the original paper, this is referred to as "spectral Adam" or "Sadam" due to its effect on the Adam optimizer update rule. The motivation behind representing the weights in the frequency domain is that optimizer updates/steps may now affect all frequencies to an equal amount. This improves the gradient conditioning, thus leading to faster convergence and increased stability at larger learning rates.
For comparison, see the TensorFlow Compression implementations of
SignalConv2DandRDFTParameter. They seem to useSignalConv2din most of their provided architectures: https://github.com/search?q=repo%3Atensorflow%2Fcompression+Conv2D&type=codeFurthermore, since this is a simple invertible transformation on the weights, it is trivial to convert any existing pretrained weights into this form via:
To override
self.weightas a property, I'm unregistering the module usingdel self._parameters["weight"]as shown in https://github.com/pytorch/pytorch/issues/46886, and also using the fact that@propertyreturns a descriptor object so thatself.weight"falls back" to the property.Checklist:
SpectralConvTransposed2dis defined correctly... Sincenn.ConvTransposed2dinherits from the same_ConvNdasnn.Conv2dwith minor adjustments, I assumed that the weights are treated in the same way.load_state_dictto be able to load/write checkpoints in a way that is compatible with the original (non-reparametrized) model definitions.torch.float32instead oftorch.complex64. Or maybe it's just the FFT which ruins things. It would be nice if there were a way to skip compilation for certain incompatible operations, particularly since we're merely taking a FFT of leaf tensors, not intermediate computations...CompressionModel.update(). Or even easier: when theSpectralConv2dmodule is inevalmode, just cache the computed weight, and invalidate/clear the cache when the module switches back totrainmode.SpectralConv2don Google, so this should be an OK name. Tensorflow Compression calls their conv layersSignalConv2d. Perhaps the most "accurate" name would beSpectralReparametrizedConv2d, but that's a bit wordy.