Closed YodaEmbedding closed 6 months ago
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[:] = k
Simple 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, :, :] = k
Simple 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[:] = k
Random 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:
++model.name="bmshj2018-factorized" ++criterion.lmbda=0.0067 ++scheduler.net.threshold=5e-4
Training forward()-loss curves comparison |
---|
Blue: Conv2d. Green: SpectralConv2d. |
Validation forward()-loss curves comparison |
---|
Blue: Conv2d. Green: SpectralConv2d. No clamp(0, 1) applied to x_hat. |
Inference/test forward()-loss curves comparison |
---|
Blue: Conv2d. Green: SpectralConv2d. No clamp(0, 1) applied to x_hat. |
However, even though the inference/test loss is evidently much smaller, the position of the point on the RD plot is actually less "optimal" w.r.t. the default configuration:
Conv2d | SpectralConv2d |
---|---|
One possibility to mitigate this might be to initially train with the spectral transform enabled for the first 30 epochs, and then disable it for the remainder of the training.
Perhaps the "inaccurate" reconstruction loss ($D$) estimation is only a problem when "noise" quantization is used during training. Not sure. Maybe modifying the covariance structure (see Balle's paper, pg 2) also makes it easier for $g_s(y + \mathcal{N}(0, 1))$ to shift the distribution and generate relatively more "optimistic" $\hat{x}$ than rounding would...? Maybe STE might help...
Contour lines of equal loss for $\lambda=0.0067$ |
---|
Still, though... a point with lower $L$ actually has worse $(R,D)$ relative to the optimal achievable $(R, D)$ for a given model architecture?!
EDIT IV (A new
hope): Found it! There's no .clamp_(0, 1)
in the forward
, which makes sense during training.
def forward(...):
x_hat = self.g_s(y_hat)
def decompress(...):
x_hat = self.g_s(y_hat).clamp_(0, 1)
By clamp_
ing the x_hat
when not in training mode, both methods now give the exact same results. It's quite surprising that simple clamping causes such a big gain in PSNR. (Between 0.66 dB to 0.74 dB.) It's also weird that the SpectralConv2d
-trained model only gets a smaller 0.66 dB jump but the Conv2d
somehow enjoys 0.74 dB. But maybe that's just luck.
EDIT V (The scheduler strikes back): According to the training losses, it looks like the SpectralConv2d
-trained model dropped LR earlier, whereas the Conv2d
-trained model happily went on for an additional 100 epochs before the first LR drop! Here were the settings I used:
scheduler:
net:
type: "ReduceLROnPlateau"
mode: "min"
factor: 0.1
patience: 10
threshold: 5e-4 # default is 1e-4
The CompressAI Trainer config is using the defaults. It's just my personal config that has the 5e-4
. Upon further reflection, it looks like that reduction in loss was slowing down anyways, so maybe scheduler.net.threshold=5e-4
is probably not that bad an idea. I'll give 1e-4
a try anyways.
EDIT VI (Return of the multi-stage training): Perhaps best is to initially train with spectral-mode, then auto-schedule switching to regular-mode some time before the first LR drop occurs.
That may (or may not?) require changing the training code, however, i.e., a new Runner
? I wonder if PyTorch has schedulers that can work with regular modules, and not just optimizers. Otherwise, I guess I could create a pseudo-optimizer that doesn't actually optimize anything, but merely switches conv.enable_transform=False
for all the convs when a pseudo-LR drop occurs... That's a convoluted way of avoiding adding a new runner. Maybe a callback is less convoluted. But then I have to write my own pseudo-LR drop logic...?!
++model.name="bmshj2018-factorized" ++criterion.lmbda=0.0067
(with default scheduler.net.threshold=1e-4
)
Inference/test forward()-loss curves comparison |
---|
Green: Conv2d. Blue: SpectralConv2d. No clamp(0, 1) applied to x_hat. |
Conv2d | SpectralConv2d |
---|---|
Much better.
Multi-stage training might still be a good idea, though, since at around loss=0.655
(i.e. 40 epochs for SpectralConv2d
) it looks like Conv2d
starts converging more quickly than SpectralConv2d
.
merged together with #270
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
SignalConv2D
andRDFTParameter
. They seem to useSignalConv2d
in 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.weight
as 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@property
returns a descriptor object so thatself.weight
"falls back" to the property.Checklist:
SpectralConvTransposed2d
is defined correctly... Sincenn.ConvTransposed2d
inherits from the same_ConvNd
asnn.Conv2d
with minor adjustments, I assumed that the weights are treated in the same way.load_state_dict
to be able to load/write checkpoints in a way that is compatible with the original (non-reparametrized) model definitions.torch.float32
instead 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 theSpectralConv2d
module is ineval
mode, just cache the computed weight, and invalidate/clear the cache when the module switches back totrain
mode.SpectralConv2d
on 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.