Package: mcquic
a.k.a.Multi-codebookQuantizers for neuralimagecompression🥳Our paper will be presented at CVPR 2022!🥳CVF Open Access|arXiv|BibTex|DemoMcQuicis a deep image compressor.Features:Solid performance and super-fast coding speed (SeeReference Models).Cross-platform support (Linux-64, Windows-64 and macOS-64, macOS-arm64).You could try the interactive demo in theHuggingFace Space!Techs:TheMcQuichold rich multi-codebooks to quantize visual features and restore images by these quantized features. Similar ideas are presented in SHA [1], VQ-VAE [2], VQ-GAN [3],etc. We summarize these as vectorized priors, and our method extends these ideas to aunified multivariate Gaussian mixture, to perform high-quality, low-latency image compression.Figure 1. Operational diagrams of different methods.Figure 2. Comparisons with traditional codecs on an image from Kodak dataset.Quick StartRequirementsConda (Recommended)DockerInstall Manually (for dev)(Optional) Install NVIDIA/ApexReference ModelsTrain a New ModelRequirementsConfigsPrepare a DatasetTrainingTestImplement MCQ by yourselfContribute to this RepositoryTo-do ListDetailed frameworkReferences and LicenseReferencesCitationCopyrightQuick StartIt is easy (with a GPU, or CPU if you like) to try our model. I would give a quick guide to help you compress an image and restore it.RequirementsTo run the model, your device needs to meet following requirements.Hardwarea CUDA-enabled GPU (≥ 8GiB VRAM, Driver version≥ 450.80.02)If you don't have GPU, running models on CPU may be slower.≥ 8GiB RAMOSI've tested all features onUbuntu, other platforms should also work. If not, pleasefile bugs.Conda (Recommended)Install this package is very easy with acondaenvironment installed,e.g.Miniconda. I recommend you to install it to a new virtual environment directly by:# Install a clean pytorch with CUDA supportcondacreate-n[ENV_NAME]python=3.9"pytorch>=1.11,<2""torchvision>=0.12,<1"cudatoolkit-cpytorch# Install mcquic and other dependenciescondainstall-n[ENV_NAME]mcquic-cxiaosu-zhu-cconda-forge
condaactivate[ENV_NAME]Above command install packages withCUDAsupport. If you just want to run it on CPU, please usecpuonlyother thancudatoolkitin the first command.Since there is no proper version of torchvision now for Apple M1, you need to change channel frompytorchtoconda-forgein the first command.Compress imagesmcquicUsage: mcquic [OPTIONS] COMMAND [ARGS]...Options:-v, --version  Print version info.-h, --help     Show this message and exit.Commands:-*        Compress/restore a file.dataset   Create training set from `images` dir to `output` dir.train     Train a model.validate  Validate a trained model from `path` by images from `images`...mcquic--helpUsage: mcquic - [OPTIONS] INPUT [OUTPUT]Compress/restore a file.Args:input (str): Input file path. If input is an image, compress it. Ifinput is a `.mcq` file, restore it.output (optional, str): Output file path or dir. If not provided, thisprogram will only print compressor information of input file.Options:-D, --debug        Set logging level to DEBUG to print verbose messages.-q, --quiet        Silence all messages, this option has higher priority to`-D/--debug`.-qp INTEGER RANGE  Quantization parameter. Higher means better image qualityand larger size.  [default: 2; 1<=x<=13]--local FILE       Use a local model path instead of download by `qp`.--disable-gpu      Use pure CPU to perform compression. This will be slow.--mse              Use model optimized for PSNR other than MsSSIM.--crop             Crop the image to align feature patches. Edges of imageare cutted though, compressed binary will be smaller.-h, --help         Show this message and exit.mcquic-qp2path/to/an/imagepath/to/output.mcqDecompress images# `-qp` is not necessary. Since this arg is written to `output.mcq`.mcquicpath/to/output.mcqpath/to/restored.pngDockerI also builddockerimagesfor you to get away from environment issues.Try with the latest docker image:dockerpullghcr.io/xiaosu-zhu/mcquic:latest# or nightly build# docker pull ghcr.io/xiaosu-zhu/mcquic:nightlyThe entrypoint of this container is set tomcquicitself. So, you can directly use it asmcquicmain program to execute.dockerrunghcr.io/xiaosu-zhu/mcquic:latest--helpTo compress/restore images, you need to mount native files into the container. Therefore, a working example forms as follows:# `someimage.png` is located in `path/to/some/folder`. And this folder will be mounted at `/workspace/workdir`.dockerrun-vpath/to/some/folder:/workspace/workdirghcr.io/xiaosu-zhu/mcquic:latest/workspace/workdir/someimage.png/workspace/workdir/output.mcq
dockerrun-vpath/to/some/folder:/workspace/workdirghcr.io/xiaosu-zhu/mcquic:latest/workspace/workdir/output.mcq/workspace/workdir/restored.pngInstall Manually (for dev)This way enables your full access to this repo for modifying. Also, if you want to go on, acondaenvironment is needed,e.g.Miniconda.Clone this repositorygitclonehttps://github.com/xiaosu-zhu/McQuic.git&&cdMcQuicCreate a virtual envmcquicand install all packages by./install.sh# for POSIX with bash.\install.ps1# for Windows with Anaconda PowerShellNow you should in themcquicvirtual environment. If not, please activate it byconda activate mcquic.Compress imagesmcquic--help
mcquic-qp2assets/sample.pngassets/compressed.mcqDecompress images# `-qp` is not necessary. Since this arg is written to `output.mcq`.mcquicassets/compressed.mcqassets/restored.pngAnd check outputs:assets/compressed.mcqandassets/restored.png.(Optional) InstallNVIDIA/ApexNVIDIA/Apexis an additional packagerequiredfor training. If you want todevelop, contribute, ortrain a new model, please ensure you've installedNVIDIA/Apexby following snippets.gitclonehttps://github.com/NVIDIA/apex&&cdapex
pipinstall-v--disable-pip-version-check--no-cache-dir--global-option="--cpp_ext"--global-option="--cuda_ext"./If you are usingDocker images, this step is not necessary.Please make sure you've installed it in the correct virtual environment.For more information such as building toolchains, please refer totheir repository.Reference ModelsI've released one pretrained model (Sorry, currently I don't have much free GPUs to train models). You could fetch them by specifying-qp [Model_NO]. Following is the pretrained model list (OthersTBA):Model No.ChannelMKThroughput (Encode/Decode)Avg.BPP------21282[8192,2048,512]25.45 Mpps / 22.03 Mpps0.1277------1219212[8192,2048,512]11.07 Mpps / 10.21 Mpps-The coding throughput is tested on a NVIDIA RTX 3090. Image file I/O, model loading,etc.are not included in the test. Throughput will be further increased by5%~15%if you convert models toTorchScript. However, it is not trivial since conversion involves entropy coder, which is a cpp extension. So, I'm not going to implement it.The main slow-down from small models to large models is caused by channel128 -> 192.Mpps = Mega-pixels per secondBPP = Bits per pixelTrain a New ModelPlease ensure you've installedNVIDIA/Apex. To train models, here are minimal and recommended system requirements.RequirementsMinimalRAM ≥ 64GiBVRAM ≥ 12GiBRecommendedVRAM ≥ 24GiBBetter if you have≥4-wayNVIDIA RTX 3090s or faster GPUs.ConfigsThe folderconfigsprovides example configexample.yamlto train models. Please refer toconfigs/README.mdfor more info.Prepare a DatasetBefore training models, you need to prepare an image dataset. It is free to pick any images to form dataset, as long as the image-size is≥512x512.To build a training dataset, please put all images in a folder (allow for sub-folders), then runmcquicdataset--helpUsage: mcquic dataset [OPTIONS] IMAGES OUTPUTCreate training set from `images` dir to `output` dir.Args:images (str): All training images folder, allow sub-folders.output (str): Output dir to create training set.Options:-D, --debug  Set logging level to DEBUG to print verbose messages.-q, --quiet  Silence all messages, this option has higher priority to`-D/--debug`.-h, --help   Show this message and exit.mcquicdatasettrain_imagesmcquic_datasetto build almdbdataset formcquicto read.Then, you could prepare a training confige.g.configs/train.yaml, and don't forget to speify dataset path.# `configs/train.yaml`...trainSet:mcquic_dataset# path to the training dataset.valSet:val_images# path to a folder of validation images.savePath:saved# path to a folder to save checkpoints....wheretrainSetandvalSetcan be any relative or absolute paths, andsavePathis a folder for saving checkpoints and logs.In this example, the final folder structure is shown below:.# A nice folder├─ 📂configs│   ...│   └── 📄train.yaml├── 📄README.md# this readme├── 📂saved# saved models apprear here├── 📂train_images# a lot of training images│   ├── 📂ImageNet│   |   ├── 📂folder1# a lot of images│   |   ├── 🖼️image1.png│   |   ...│   ├── 📂COCO│   |   ├── 🖼️image1.png│   |   ├── 🖼️image2.png│   |   ...|   ...├── 📂mcquic_dataset# generated training dataset|├── 📀data.mdb|├── 📀lock.mdb|└── 📄metadata.json└── 📂val_images# a lot of validation images├── 🖼️image1.png├── 🖼️image2.png...TrainingTo train a new model, runmcquictrain--helpUsage: mcquic train [OPTIONS] [CONFIG]Train a model.Args:config (str): Config file (yaml) path. If `-r/--resume` is present butconfig is still given, then this config will be used to update theresumed training.Options:-D, --debug        Set logging level to DEBUG to print verbose messages.-q, --quiet        Silence all messages, this option has higher priority to`-D/--debug`.-r, --resume FILE  `.ckpt` file path to resume training.-h, --help         Show this message and exit.mcquictrainconfigs/train.yamland saved model is located insaved/mcquic_dataset/latest.To resume an interuptted training, runmcquictrain-r, ormcquictrain-rconfigs/train.yamlif you want to use an updated config (e.g. tuned learning rate, modified hyper-parameters) to resume training.TestYou could use any save checkpoints (usually located in abovesavePath) to validate the performance. For examplemcquicvalidate--helpUsage: python -m mcquic.validate [OPTIONS] PATH IMAGES [OUTPUT]Validate a trained model from `path` by images from `images` dir, andpublish a final state_dict to `output` path.Args:path (str): Saved checkpoint path.images (str): Validation images folder.output (str): Dir to save all restored images.Options:-D, --debug        Set logging level to DEBUG to print verbose messages.-q, --quiet        Silence all messages, this option has higher priority to`-D/--debug`.-e, --export PATH  Path to export the final model that is compatible withmain program.-h, --help         Show this message and exit.mcquicvalidate-epath/to/final/modelpath/to/a/checkpointpath/to/images/folderpath/to/output/folderAnd the output "final/model" is compatible with the main programmcquic, you could directly use this local model to perform compression. Try:mcquic--localpath/to/final/modelassets/sample.pngassets/compressed.mcq# `--local` is not necessary. Since this arg is written to `output.mcq`.mcquicassets/compressed.mcqassets/restored.pngIf you think your model is awesome, please don't hasitate toContribute to this Repository!Implement MCQ by yourselfA minimal implementation of the multi-codebook quantizer comes up with (please refer toquantizer.pyfor notes):classQuantizer(nn.Module):"""Quantizer with `m` sub-codebooks,`k` codewords for each, and`n` total channels.Args:m (int): Number of sub-codebooks.k (int): Number of codewords for each sub-codebook.n (int): Number of channels of latent variables."""def__init__(self,m:int,k:int,n:int):super().__init__()# A codebook, feature dim `d = n // m`.self._codebook=nn.Parameter(torch.empty(m,k,n//m))self._initParameters()def_initParameters(self):nn.init.normal_(self._codebook,std=math.sqrt(2/(5*n/m)))defforward(self,x:Tensor,t:float=1.0)->(Tensor,Tensor):"""Module forward.Args:x (Tensor): Latent variable with shape [b, n, h, w].t (float, 1.0): Temperature for Gumbel softmax.Return:Tensor: Quantized latent with shape [b, n, h, w].Tensor: Binary codes with shape [b, m, h, w]."""b,_,h,w=x.shape# [b, m, d, h, w]x=x.reshape(b,len(self._codebook),-1,h,w)# [b, m, 1, h, w], square of xx2=(x**2).sum(2,keepdim=True)# [m, k, 1, 1], square of codebookc2=(self._codebook**2).sum(-1,keepdim=True)[...,None]# [b, m, d, h, w] * [m, k, d] -sum-> [b, m, k, h, w], dot product between x and codebookinter=torch.einsum("bmdhw,mkd->bmkhw",x,self._codebook)# [b, m, k, h, w], pairwise L2-distancedistance=x2+c2-2*inter# [b, m, k, h, w], distance as logits to samplesample=F.gumbel_softmax(-distance,t,hard=True,dim=2)# [b, m, d, h, w], use sample to find codewordsquantized=torch.einsum("bmkhw,mkd->bmdhw",sample,self._codebook)# back to [b, n, h, w]quantized=quantized.reshape(b,-1,h,w)# [b, n, h, w], [b, m, h, w], quantizeds and binariesreturnquantized,sample.argmax(2)Contribute to this RepositoryIt will be very nice if you want to check your new ideas or add new functions 😊. You will need to installmcquicbyDockerormanually (with optional step). Just like other git repos, before raising issues or pull requests, please take a thorough look atissue templates.To-do Listmcquic serviceMore pretrained modelDetailed frameworkThanks for your attention!❤️ Here are details in the paper.Following previous works, we build the compression model as an AutoEncoder. Bottleneck of encoder (analysis transform) outputs a small feature map and is quantized bymulti-codebook vector-quantizationother than scalar-quantization. Quantizers are cascaded to effectively estimate latent distribution.Figure 3. Left: Overall framework. Right: Structure of a quantizer.Right part of above figure shows detailed structure of our proposed quantizer.References and LicenseReferences[1] Agustsson, Eirikur, et al. "Soft-to-hard vector quantization for end-to-end learning compressible representations." NeurIPS 2017.[2] Van Den Oord, Aaron, and Oriol Vinyals. "Neural discrete representation learning." NeurIPS 2017.[3] Esser, Patrick, Robin Rombach, and Bjorn Ommer. "Taming transformers for high-resolution image synthesis." CVPR 2021.CitationTo cite our paper, please use following BibTex:@inproceedings{McQuic,
  author    = {Xiaosu Zhu and
               Jingkuan Song and
               Lianli Gao and
               Feng Zheng and
               Heng Tao Shen},
  title     = {Unified Multivariate Gaussian Mixture for Efficient Neural Image Compression},
  booktitle = {CVPR},
  % pages     = {????--????}
  year      = {2022}
}CopyrightFonts:Source Sans Pro. © 2010, 2012 Adobe Systems Incorporated, SIL Open Font License.Flash Rogers 3D. © 2007 Iconian Fonts, donationware.Cambria Math. © 2017 Microsoft Corporation. All rights reserved.Times New Roman. © 2017 The Monotype Corporation. All Rights Reserved.Caramel and Vanilla. © 2017 FOUND MY FONT LTD. All Rights Reserved.Pictures:kodim24.pngby Alfons Rudolph, Kodak Image Dataset.assets/sample.pngby Ales Krivec, CLIC Professional valid set.Third-party repos:ReposLicensePyTorchBSD-styleTorchvisionBSD-3-ClauseApexBSD-3-ClauseTensorboardApache-2.0KorniaApache-2.0richMITpython-lmdbOpenLDAP Version 2.8PyYAMLMITmarshmallowMITclickBSD-3-ClausevlutilsApache-2.0MessagePackApache-2.0pybind11BSD-styleCompressAIBSD 3-Clause ClearTaming-transformerMITmarshmallow-jsonschemaMITjson-schema-for-humansApache-2.0CyclicLRMITbatch-transformsMITpytorch-msssimMITStreamlitApache-2.0condaBSD 3-ClauseThis repo is licensed underApache LicenseVersion 2.0