VQGAN-f16-16384

模型描述

这是VQGAN的Flax/JAX实现，通过利用卷积方法和transformer来学习一个上下文丰富的可视部分的码书。它在 Taming Transformers for High-Resolution Image Synthesis 中被引入（ CVPR paper ）。

该模型允许将图像编码为从码书中提取的固定长度的标记序列。

该版本的模型使用了一个降低因子f=16和13384个标记的词汇表。

作为降低因子工作原理的示例，大小为256x256的图像被编码为256个标记的序列：256/16 * 256/16。大小为512x512的图像将得到1024个标记的序列。

用于训练的数据集

• ImageNet。我们没有从头开始训练该模型。而是从 a checkpoint pre-trained on ImageNet 开始。
• Conceptual Captions 3M （CC3M）。
• OpenAI subset of YFCC100M 。

我们在CC3M和YFCC100M上进行了微调，以提高人物和面部的编码质量，这在ImageNet中表示得不够好。我们在CC3M和YFCC100M的子集中使用了2268720张图像来完成此任务。

训练过程

在PyTorch中使用 taming-transformers 进行微调。完整的训练过程和模型准备包括以下步骤：

• 在ImageNet上进行预训练。之前已经完成。我们使用了 this checkpoint 。
• 微调， Part 1 。
• 微调， Part 2 -继续第一部分。最终的检查点已上传到 boris/vqgan_f16_16384 。
• 转换为JAX，这是本文中描述的模型。

如何使用

可以使用VQModel的 Suraj Patil's implementation 来加载检查点。

• 基于 Suraj 的工作的示例笔记本：

• 使用JAX的批处理编码，包括使用PyTorch进行数据加载的完整示例：

# VQGAN-JAX - pmap encoding HowTo

import numpy as np

# For data loading
import torch
import torchvision.transforms.functional as TF
from torch.utils.data import Dataset, DataLoader
from torchvision.datasets.folder import default_loader
from torchvision.transforms import InterpolationMode

# For data saving
from pathlib import Path
import pandas as pd
from tqdm import tqdm

import jax
from jax import pmap

from vqgan_jax.modeling_flax_vqgan import VQModel

## Params and arguments

# List of paths containing images to encode
image_list = '/sddata/dalle-mini/CC12M/10k.tsv'
output_tsv = 'output.tsv'                # Encoded results
batch_size = 64
num_workers = 4    # TPU v3-8s have 96 cores, so feel free to increase this number when necessary

# Load model
model = VQModel.from_pretrained("flax-community/vqgan_f16_16384")

## Data Loading.

# Simple torch Dataset to load images from paths.
# You can use your own pipeline instead.
class ImageDataset(Dataset):
    def __init__(self, image_list_path: str, image_size: int, max_items=None):
        """
        :param image_list_path: Path to a file containing a list of all images. We assume absolute paths for now.
        :param image_size: Image size. Source images will be resized and center-cropped.
        :max_items: Limit dataset size for debugging
        """
        self.image_list = pd.read_csv(image_list_path, sep='\t', header=None)
        if max_items is not None: self.image_list = self.image_list[:max_items]
        self.image_size = image_size
        
    def __len__(self):
        return len(self.image_list)
    
    def _get_raw_image(self, i):
        image_path = Path(self.image_list.iloc[i][0])
        return default_loader(image_path)
    
    def resize_image(self, image):
        s = min(image.size)
        r = self.image_size / s
        s = (round(r * image.size[1]), round(r * image.size[0]))
        image = TF.resize(image, s, interpolation=InterpolationMode.LANCZOS)
        image = TF.center_crop(image, output_size = 2 * [self.image_size])
        image = np.expand_dims(np.array(image), axis=0)
        return image
    
    def __getitem__(self, i):
        image = self._get_raw_image(i)
        return self.resize_image(image)

## Encoding

# Encoding function to be parallelized with `pmap`
# Note: images have to be square
def encode(model, batch):
    _, indices = model.encode(batch)
    return indices

# Alternative: create a batch with num_tpus*batch_size and use `shard` to distribute.
def superbatch_generator(dataloader, num_tpus):
    iter_loader = iter(dataloader)
    for batch in iter_loader:
        superbatch = [batch.squeeze(1)]
        try:
            for _ in range(num_tpus-1):
                batch = next(iter_loader)
                if batch is None:
                    break
                # Skip incomplete last batch
                if batch.shape[0] == dataloader.batch_size:
                    superbatch.append(batch.squeeze(1))
        except StopIteration:
            pass
        superbatch = torch.stack(superbatch, axis=0)
        yield superbatch

def encode_dataset(dataset, batch_size=32):
    dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)
    superbatches = superbatch_generator(dataloader, num_tpus=jax.device_count())
    
    num_tpus = jax.device_count()
    dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)
    superbatches = superbatch_generator(dataloader, num_tpus=num_tpus)
    
    p_encoder = pmap(lambda batch: encode(model, batch))

    # Save each superbatch to avoid reallocation of buffers as we process them.
    # Keep the file open to prevent excessive file seeks.
    with open(output_tsv, "w") as file:
        iterations = len(dataset) // (batch_size * num_tpus)
        for n in tqdm(range(iterations)):
            superbatch = next(superbatches)
            encoded = p_encoder(superbatch.numpy())
            encoded = encoded.reshape(-1, encoded.shape[-1])

            # Extract paths from the dataset, save paths and encodings (as string)
            start_index = n * batch_size * num_tpus
            end_index = (n+1) * batch_size * num_tpus
            paths = dataset.image_list[start_index:end_index][0].values
            encoded_as_string = list(map(lambda item: np.array2string(item, separator=',', max_line_width=50000, formatter={'int':lambda x: str(x)}), encoded))
            batch_df = pd.DataFrame.from_dict({"image_file": paths, "encoding": encoded_as_string})
            batch_df.to_csv(file, sep='\t', header=(n==0), index=None)
            
dataset = ImageDataset(image_list, image_size=256)
encoded_dataset = encode_dataset(dataset, batch_size=batch_size)

Hub中的相关模型

• 在同样描述的数据集上训练的VQGAN的PyTorch版本： boris/vqgan_f16_16384 。
• DALL·E mini ，OpenAI的DALL·E的Flax/JAX简化实现。

其他

该模型已成功用于 DALL·E mini 的实现的一部分。我们的 report 包含了如何在图像编码/生成流水线中利用它的更多细节。