FFCV là một hệ thống tải dữ liệu trực tiếp giúp tăng đáng kể tốc độ truyền dữ liệu trong quá trình huấn luyện mô hình:

• Đào tạo mô hình ImageNet trên một GPU mất 35 phút (chi phí khoảng 98 cent mỗi mô hình trên AWS)
• Đào tạo mô hình CIFAR-10 trên một GPU mất 36 giây (chi phí khoảng 2 cent mỗi mô hình trên AWS)

Cài Đặt Môi Trường

Theo tài liệu chính thức: https://docs.ffcv.io/

1. conda-forge

conda create -n ffcv python=3.10 pkg-config opencv -c conda-forge --override-channels
conda activate ffcv

Lưu ý: pkg-config giúp tìm thấy opencv4.pc, cần thiết để biên dịch ffcv.

2. torch

conda install pytorch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 pytorch-cuda=11.8 -c pytorch -c nvidia

3. ffcv

pip install ffcv

Quy Trình Chi Tiết

Dưới đây là một số bước tôi đã tổng hợp từ tài liệu chính thức.

1. Lưu trữ Dữ Liệu dưới Định dạng FFCV

from ffcv.writer import DatasetWriter
from ffcv.fields import RGBImageField
from torchvision import transforms
from PIL import Image
from torch.utils.data import Dataset
import os

class ImagePairSet(Dataset):
    def __init__(self, source_folder, target_folder, transform=None):
        self.source_folder = source_folder
        self.target_folder = target_folder
        self.transform = transform
        self.file_names = sorted(os.listdir(source_folder))

    def __len__(self):
        return len(self.file_names)

    def __getitem__(self, index):
        source_path = os.path.join(self.source_folder, self.file_names[index])
        target_path = os.path.join(self.target_folder, self.file_names[index])
        source_image = Image.open(source_path).convert('RGB')
        target_image = Image.open(target_path).convert('RGB')

        if self.transform:
            source_image = self.transform(source_image)
            target_image = self.transform(target_image)

        return source_image, target_image

source_folder = 'xx/source/'
target_folder = 'xx/target/'

dataset = ImagePairSet(source_folder, target_folder, transform=None)
writer = DatasetWriter('data.beton', {
    'source': RGBImageField(write_mode='smart'),
    'target': RGBImageField(write_mode='smart')
})
writer.from_indexed_dataset(dataset)

2. Kiểm Tra Việc Chuyển Đổi

from ffcv.loader import Loader, OrderOption

loader = Loader('data.beton', batch_size=32, num_workers=4, order=OrderOption.RANDOM)

for batch in loader:
    print(type(batch))
    source_images, target_images = batch[0], batch[1]
    print(source_images.shape, target_images.shape)
    break

3. Thay Thế DataLoader của PyTorch

training_loader = Loader(
    dataset_config["beton_file"],
    batch_size=args.batch_size,
    num_workers=4,
    order=OrderOption.RANDOM,
    pipelines={
        'source': [SimpleRGBImageDecoder(), ToTensor(), ToTorchImage(), Convert(torch.float32),
                   torchvision.transforms.Normalize(CIFAR_MEAN, CIFAR_STD), ToDevice(torch.device('cuda:0'))],
        'target': [SimpleRGBImageDecoder(), ToTensor(), ToTorchImage(), Convert(torch.float32),
                   torchvision.transforms.Normalize(CIFAR_MEAN, CIFAR_STD), ToDevice(torch.device('cuda:0'))],
    },
    drop_last=True
)

4. Tạo Các Thao Tác Pipeline Tùy Chỉnh

Ngoài các thao tác được cung cấp sẵn, đôi khi chúng ta cần triển khai các thao tác tiền xử lý và tăng cường dữ liệu riêng.

• Khai báo phương thức `declare_state_and_memory` sau mỗi hàm
• Kế thừa từ `Operation` trong `ffcv.pipeline.operation`

Ví dụ đơn giản về một thao tác chuẩn hóa:

class NormalizeFirstChannel(Operation):
    def __init__(self, mean_value, std_value):
        super().__init__()
        self.mean = mean_value
        self.std = std_value

    def generate_code(self):
        mean = self.mean
        std = self.std
        def normalize(images, *_):
            B, C, H, W = images.shape
            for i in range(B):
                images[i, 0, :, :] = (images[i, 0, :, :] - mean) / std
            return images

        return normalize

    def declare_state_and_memory(self, previous_state: State):
        return previous_state, 0