vector-quantize-pytorch

收录理由

向量量化能把连续数据变成离散编码，这个库用一组干净的 PyTorch 模块把相关方法打包好了。DeepMind 的 VQ-VAE-2 生成图像、OpenAI 的 Jukebox 生成音乐，底层用的就是这套技术；在这里你可以直接用 VectorQuantize、ResidualVQ 和 GroupedResidualVQ 这些类，把它们嵌进自己的模型。默认情况下，码本通过指数移动平均来更新，也可以换成 DiVeQ 方式——靠梯度学习码本，不需要额外的辅助损失。另外还支持 k-means 初始化和随机码采样这些附加功能。做 VAE 或者基于 token 的图像、音频、视频模型的研发人员，能省下不少重复实现近期文献的功夫；不过要留意，它提供的是专注的构建模块，不是端到端的生成工具。