vector-quantize-pytorch

项目介绍：Vector Quantize Pytorch

项目概述

Vector Quantize Pytorch 是一个基于向量量化的库，它最初从 Deepmind 的 TensorFlow 实现中转录而来，如今被打包成了一个方便使用的 Python 包。这个库主要利用指数滑动平均来更新字典，已被 Deepmind 和 OpenAI 成功运用在高质量的图像生成（如 VQ-VAE-2）和音乐生成（如 Jukebox）中。

项目安装

你只需执行以下命令即可轻松安装：

$ pip install vector-quantize-pytorch

主要功能

向量量化（Vector Quantize）

这个功能允许用户创建一个向量量化实例，指定向量的维度、码本的大小、指数衰减率以及承诺损失权重。示例代码如下：

import torch
from vector_quantize_pytorch import VectorQuantize

vq = VectorQuantize(
    dim = 256,
    codebook_size = 512,
    decay = 0.8,
    commitment_weight = 1.
)

x = torch.randn(1, 1024, 256)
quantized, indices, commit_loss = vq(x)

残差向量量化（Residual VQ）

这一功能可以通过递归量化波形的残差来提高量化效果。通过使用 ResidualVQ 类，你可以指定多量化器的个数以及码本的大小：

import torch
from vector_quantize_pytorch import ResidualVQ

residual_vq = ResidualVQ(
    dim = 256,
    num_quantizers = 8,
    codebook_size = 1024
)

x = torch.randn(1, 1024, 256)
quantized, indices, commit_loss = residual_vq(x)

进阶技巧

增强码本使用率

降低码本维度：通过调整编码器的输入和输出维度，可以提升码本使用率。
余弦相似度：用余弦相似度代替传统的距离度量，有助于提高码本使用率，并改善重构效果。
淘汰陈旧码：设定一个阈值，自动淘汰低于使用次数的码字。

增加正交正则化损失

通过在图像等应用中对码本施加正交性的约束，让离散码具备翻译等效性，可以显著提升后续任务的效果：

import torch
from vector_quantize_pytorch import VectorQuantize

vq = VectorQuantize(
    dim = 256,
    codebook_size = 256,
    orthogonal_reg_weight = 10
)

img_fmap = torch.randn(1, 256, 32, 32)
quantized, indices, loss = vq(img_fmap)

随机投影量化（Random Projection Quantizer）

该方法不需要学习一个量化器，而是通过一个随机初始化的矩阵将信号投影，并匹配到随机初始化的码本。这项技术已经在大型的语音模型中得到应用。

import torch
from vector_quantize_pytorch import RandomProjectionQuantizer

quantizer = RandomProjectionQuantizer(
    dim = 512,
    num_codebooks = 16,
    codebook_dim = 256,
    codebook_size = 1024
)

x = torch.randn(1, 1024, 512)
indices = quantizer(x)