Optimizing PyTorch Data Pipelines: From Bottlenecks to 39× Speedups

Giới thiệu#

Hôm nay tôi thử nghiên cứu lại quy trình huấn luyện PyTorch cơ bản trước khi đi vào thử nghiệm một số thuật toán computer vision mới. Mục tiêu chính là để kiểm tra xem những hiểu biết và thói quen trước đó của tôi liệu có còn hiệu quả, hay còn thiếu sót gì trong bối cảnh hiện nay hay không. Nói thẳng ra là, đôi khi chúng ta dùng các framework "đóng gói" (như PyTorch Lightning, Hugging Face Accelerate) nhiều quá mà quên mất những gì thực sự xảy ra bên dưới.

Để thực hiện các thử nghiệm này, tôi sử dụng một máy workstation cá nhân với cấu hình khá mạnh, được thiết kế cho các tác vụ tính toán và AI:

• CPU: Intel i9-14900K (32 cores) @ 5.700GHz
• GPU: NVIDIA RTX A4000 (16GB VRAM)
• RAM: 32 GB

Một cấu hình mạnh như thế này, đặc biệt là GPU, sẽ khiến bất kỳ sự chậm chạp nào từ phía CPU hoặc pipeline I/O (tải dữ liệu) trở nên rõ ràng hơn. Nếu GPU phải chờ CPU chuẩn bị dữ liệu, chúng ta sẽ thấy ngay hiện tượng "GPU starvation" (GPU bị bỏ đói).

Về cơ bản, một training loop tiêu chuẩn cho bài toán học có giám sát (supervised learning) trong PyTorch khá đơn giản. Nó được minh họa qua đoạn code dưới đây:

# Ví dụ cho một thuật supervised learning có label
for epoch in range(10):
    for batch_idx, (x, y) in enumerate(train_loader):
        # Mỗi vòng lặp xử lý một mini-batch:
        #   x: tensor đầu vào [batch_size, ...]
        #   y: labels cho mini-batch
        #   batch_idx: chỉ số của mini-batch
 
        # 1. Reset gradients từ vòng lặp trước
        optimizer.zero_grad()                    
        
        # 2. Forward pass: Đưa dữ liệu qua model
        predictions = model(x)                    
        
        # 3. Tính toán loss
        loss = loss_function(predictions, y)      
        
        # 4. Backward pass: Tính toán gradient
        loss.backward()                           
        
        # 5. Cập nhật trọng số
        optimizer.step()                          

Experimental datasets#

Trong phần thử nghiệm, tôi sử dụng hai bộ dữ liệu kinh điển trong computer vision: MNIST và FashionMNIST.

• MNIST: Tập dữ liệu chữ số viết tay do Yann LeCun công bố, gồm 10 lớp tương ứng với các chữ số từ 0 đến 9.
• FashionMNIST: Được đề xuất như một phiên bản thay thế mang tính thử thách hơn cho MNIST, gồm hình ảnh của các sản phẩm thời trang như áo thun, quần dài, giày, túi xách,…

Để minh họa trực quan, Hình 1 cho thấy một số mẫu ảnh điển hình từ FashionMNIST, trong khi Hình 2 trình bày các mẫu ảnh tương ứng từ MNIST. Mỗi hình ảnh đều mang định dạng grayscale kích thước 28×28, thể hiện độ đơn giản nhưng đủ thông tin để phục vụ cho các bài toán phân loại cơ bản.

Cả hai bộ dữ liệu đều có đặc điểm chung:

• 60.000 ảnh dùng cho huấn luyện và 10.000 ảnh dành cho kiểm thử.
• Hình ảnh dạng grayscale kích thước 28×28.
• Mục tiêu là nhận diện và phân loại ảnh vào đúng 1 trong 10 nhãn có sẵn.

Nhờ kích thước nhỏ gọn và dễ xử lý, MNIST và FashionMNIST thường được xem như “Hello World” của bài toán Computer Vision. Chúng cho phép chúng ta tập trung vào tối ưu hóa pipeline tính toán và cơ chế nạp dữ liệu, mà không bị chi phối bởi các biến đổi dữ liệu phức tạp hay chi phí I/O lớn từ ổ đĩa.

Một khám phá thú vị: CSV vs. Parquet#

Lúc này tôi nhận thấy các file .csv mà tôi đang lưu trữ (mỗi hàng là 784 pixel + 1 label) khá là nặng. Tôi chợt nhớ đến hồi xưa có đọc đâu đó về Parquet và thử lưu trữ dataset với format này xem sao.

Ồ và kết quả thật thú vị: dung lượng nhỏ hơn đáng kể và tốc độ đọc ghi nhanh hơn hẳn.

Lý do đằng sau sự hiệu quả này:

• Lưu trữ theo cột (Columnar Storage): Parquet lưu trữ theo cột, còn CSV lưu trữ theo hàng. Khi dữ liệu được lưu theo cột, các giá trị trong cùng một cột (ví dụ: pixel_10, pixel_11...) thường có cùng kiểu dữ liệu và tính chất tương đồng. Điều này giúp nén tốt hơn rất nhiều.
• Nén và Encoding hiệu quả: Parquet hỗ trợ nhiều codec nén (Snappy, Gzip, Zstandard…) và các kỹ thuật encoding đặc biệt cho từng cột. Vì dữ liệu cùng kiểu trong cột dễ tìm thấy pattern hơn, việc nén hiệu quả hơn đáng kể.

Nhược điểm: Parquet là định dạng nhị phân. Bạn không thể mở nó bằng text editor và đọc như file CSV được. Nhưng với data lớn, đây là một sự đánh đổi hoàn toàn xứng đáng.

Naive Implementation#

Ok, lúc này chúng ta đi vào implement thử. Chúng ta cần một class Dataset tùy chỉnh để đọc file Parquet và một model CNN đơn giản. Tôi cố định random_seed = 42 để đảm bảo kết quả nhất quán.

Class Dataset (Version 1)#

Đây là cách triển khai "ngây thơ" (naive) đầu tiên. Ý tưởng là load toàn bộ file Parquet vào RAM trong __init__, sau đó hàm __getitem__ sẽ chịu trách nhiệm lấy ra hàng (.iloc), xử lý và biến đổi nó thành Tensor.

class MNISTDataset(Dataset):
    def __init__(self, parquet_path: str, num_classes: int = 10):
        # Load toàn bộ file vào RAM
        self.df = pd.read_parquet(parquet_path)
        self.label_col = 'label' if 'label' in self.df.columns else None
        self.feature_cols = [c for c in self.df.columns if c != self.label_col]
        self.num_classes = num_classes
 
    def __len__(self):
        return len(self.df)
    
    def __getitem__(self, idx):
        # 1. Lấy dữ liệu theo hàng
        row = self.df.iloc[idx]
        
        # 2. Xử lý features
        pixels = row[self.feature_cols].to_numpy(dtype=np.float32)
        image = torch.tensor(pixels.reshape(1, 28, 28)) / 255.0
 
        # 3. Xử lý label
        if self.label_col:
            # Trả về label dạng số nguyên (index)
            label = torch.tensor(int(row[self.label_col]), dtype=torch.long)
        else:
            label = torch.tensor(-1, dtype=torch.long) # Cho trường hợp test (không có label)
 
        return {
            "image": image,          # Tensor [1, 28, 28]
            "label": label,          # Tensor [1]
        }

Tiny Neural Network#

Mô hình này chỉ là một CNN cơ bản để có thứ gì đó cho GPU tính toán.

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, 3, padding=1)   # [B, 1, 28, 28] → [B, 16, 28, 28]
        self.pool  = nn.MaxPool2d(2, 2)               # [B, 16, 28, 28] → [B, 16, 14, 14]
        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)  # [B, 16, 14, 14] → [B, 32, 14, 14]
        self.fc1   = nn.Linear(32 * 7 * 7, 128)
        self.fc2   = nn.Linear(128, num_classes)
 
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(x.size(0), -1)  # flatten
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

Model này có số lượng tham số là 206,922 (một tiny network).

Hyperparameters#

Tôi thử tập FashionMNIST trước, config như sau:

train_dataset = MNISTDataset("data/FashionMNIST/train.parquet")
test_dataset = MNISTDataset("data/FashionMNIST/test.parquet")
 
train_dataloader = DataLoader(
    train_dataset,
    batch_size=16,
    shuffle=True,
)
 
test_dataloader = DataLoader(
    test_dataset,
    batch_size=16,
    shuffle=False,
)
 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
model = SimpleCNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

Tôi bắt đầu train thử 5 epochs.

Lưu ý: Đáng lẽ tôi phải chia tập train ra thành train và val để đánh giá trên tập val, chứ không phải tập test. Nhưng tạm thời tôi chưa quan tâm đến nó vội. Ngoài ra, việc gọi model.train() và model.eval() rất quan trọng (đặc biệt khi dùng Dropout hay BatchNorm), nhưng với model đơn giản này, tôi tạm bỏ qua.

Ok và log của batch_size=16:

Average Train Time per Epoch: 17.13s
Average Test Time per Epoch:  2.39s
Total Train Time: 85.64s
Total Test Time:  11.95s

Ừm như vậy là một epoch mất trung bình khoảng 19 giây (17s train + 2s test).

Phân tích Bottleneck trong Pipeline#

Một điều đến rất tự nhiên là: tăng batch size lên để tăng tốc độ tính toán - như tôi vẫn thường làm. GPU RTX A4000 của tôi đang quá nhàn rỗi với batch_size=16. Tôi thử tăng batch_size lên 32, 64, 128, 256, 512.

Lúc này một điều kỳ lạ đã xảy ra. Quả thật tốc độ đã tăng lên một chút (trung bình khoảng 16 giây/epoch), nhưng kể từ batch_size=32 cho đến tận batch_size=512 thì thời gian không hề thay đổi đáng kể.

Dưới đây là log của batch_size=512:

Average Train Time per Epoch: 13.67s
Average Test Time per Epoch:  2.17s
Total Train Time: 68.36s
Total Test Time:  10.87s

Nó chỉ nhanh hơn một chút không đáng kể. Chuyện quái gì đang diễn ra vậy? - tôi nghĩ trong đầu.

Sau đó tôi nhận ra ngay vấn đề, cũng phải cảm ơn lúc học về CUDA mà phản ứng lúc này của tôi tốt hơn hẳn lúc xưa. Tôi đoán ra ngay là vấn đề không nằm ở compute (GPU) mà nằm ở data pipeline.

Cụ thể hơn là: load data không đủ nhanh để theo kịp với tốc độ tính toán của GPU. Một bottleneck kinh điển. GPU utilization lúc này rất thấp, nó dành phần lớn thời gian chờ CPU đưa batch mới.

Ok, tôi bắt đầu kiểm tra kỹ lên các đoạn code trên. Tôi nhận ra có hai điểm có tiềm năng cải thiện:

• Các thuộc tính của DataLoader.
• Một bottleneck nghiêm trọng nằm ở bên trong hàm __getitem__ của class Dataset.

Lưu ý: Class Dataset kia của tôi implement trước đó là load toàn bộ data (file Parquet) vào RAM (lưu trong self.df). Trên thực tế nếu dữ liệu lớn (hàng trăm GB) thì không thể nào làm như vậy được, lúc đó phải chia nhỏ file hoặc streaming data. Nhưng trong trường hợp này, data đủ nhỏ để fit vào RAM.

Tối ưu DataLoader để Tăng Throughput#

Tôi bắt đầu thử nghiệm với DataLoader trước, giữ nguyên batch_size=512.

num_workers#

num_workers là số tiến trình (process) con mà DataLoader sẽ spawn để load dữ liệu song song. Nếu không đề cập gì (mặc định là 0), dữ liệu sẽ được load trong main process. Đây chính là bottleneck! Main process đang bận điều phối GPU, lại phải làm thêm việc tải dữ liệu. Như vậy về mặt lý thuyết, nếu chúng ta tăng số lượng worker lên thì sẽ tăng tốc độ. Tôi sẽ chạy thử với num_workers=4 (CPU của tôi có 32 cores, nhưng 4 là con số khởi đầu hợp lý).

train_dataloader = DataLoader(
    train_dataset,
    batch_size=512,
    shuffle=True,
    num_workers=4
)
 
test_dataloader = DataLoader(
    test_dataset,
    batch_size=512,
    shuffle=False,
    num_workers=4
)

Wow! Và kết quả thật đáng kinh ngạc. Trung bình mỗi epoch bây giờ chỉ khoảng 4.3 giây, giảm khoảng gần 4 lần so với trước đó (16 giây).

Average Train Time per Epoch: 3.68s
Average Test Time per Epoch:  0.67s
Total Train Time: 18.39s
Total Test Time:  3.37s

Ok, lúc này chúng ta đặt ra vấn đề là sẽ ra sao nếu như tăng số lượng worker (ví dụ 8, 16)? Thực tế thì Intel i9-14900K có tận 32 cores, khi tôi thử nghiệm tăng worker lên thì hiện tượng lại giống như batch size: tốc độ không giảm đáng kể cho lắm. Có vẻ num_workers=4 đã đủ để bão hòa pipeline. Tôi cố định config là num_workers=4.

persistent_workers#

Trong DataLoader thì tham số này mặc định là False. Khi đó, DataLoader hoạt động như sau:

• Bắt đầu epoch $\rightarrow$ tạo num_workers process.
• Workers load batch.
• Kết thúc epoch $\rightarrow$ kill toàn bộ workers.
• Sang epoch tiếp theo $\rightarrow$ spawn lại workers từ đầu.

Overhead ở đây là: thời gian tạo process mới, reload lại file dataset, reload library (numpy/pandas...).

Nếu persistent_workers=True, các worker sẽ được giữ sống qua các epoch. Hiệu quả chủ yếu nhìn thấy rõ đối với số lượng epoch lớn, nhưng mà tôi vẫn cứ chạy với 5 epoch xem sao.

Average Train Time per Epoch: 3.61s
Average Test Time per Epoch:  0.62s
Total Train Time: 18.06s
Total Test Time:  3.09s

Hmm, cũng không đáng kể cho lắm.

pin_memory#

Tiếp theo chúng ta sẽ nói về pin_memory. Thuộc tính này rất quan trọng, liên quan trực tiếp đến cách PyTorch chuyển dữ liệu từ RAM (CPU) sang VRAM (GPU) sao cho nhanh và hiệu quả hơn.

Để hiểu pin_memory=True làm gì, chúng ta cần biết về hai loại bộ nhớ của CPU:

• Bộ nhớ thông thường (Pageable Memory): Khi chúng ta tạo một tensor hoặc một biến bất kỳ, nó mặc định được lưu trữ trong pageable memory. Đây là bộ nhớ mà Hệ điều hành (OS) có toàn quyền quản lý nó. Nếu OS thấy RAM bị thiếu, nó có thể lấy khối dữ liệu đó swap ra ổ cứng để giải phóng RAM cho các tác vụ khác.

• Bộ nhớ "Ghim" (Page-Locked / Pinned Memory): Đây là một loại bộ nhớ đặc biệt trong RAM của CPU. Khi chúng ta pin một vùng nhớ, chúng ta đang nói với OS là: "Tuyệt đối không được di chuyển hay swap vùng dữ liệu này ra ổ cứng." Nó được khóa tại một địa chỉ vật lý cố định trong RAM.

Vấn đề nằm ở đâu?

Việc truyền dữ liệu từ CPU sang GPU sử dụng cơ chế DMA (Direct Memory Access) để đạt tốc độ cao nhất. DMA yêu cầu dữ liệu phải có một địa chỉ vật lý cố định trong RAM. Nó không thể làm việc hiệu quả với Pageable Memory vì OS có thể di chuyển dữ liệu đó đi bất cứ lúc nào, buộc GPU phải chờ hoặc thực hiện một bước copy trung gian.

Khi chúng ta set pin_memory=True trong DataLoader, nó sẽ tự động tải các batch dữ liệu vào Pinned Memory (thay vì Pageable Memory). Vì vùng nhớ này là cố định, cơ chế DMA của GPU có thể truy cập và sao chép nó trực tiếp sang VRAM, loại bỏ độ trễ và tăng tốc độ truyền tải đáng kể.

Ok, lý thuyết là vậy, tôi thử bật nó lên (cùng với num_workers=4 và persistent_workers=True):

train_dataloader = DataLoader(
    train_dataset,
    batch_size=512,
    shuffle=True,
    num_workers=4,
    persistent_workers=True,
    pin_memory=True
)
 
test_dataloader = DataLoader(
    test_dataset,
    batch_size=512,
    shuffle=False,
    num_workers=4,
    persistent_workers=True,
    pin_memory=True
)

Và kết quả:

Average Train Time per Epoch: 3.64s
Average Test Time per Epoch:  0.62s
Total Train Time: 18.19s
Total Test Time:  3.09s

Rất tiếc là lý thuyết rất hoành tráng nhưng nó hầu như không có tác động gì nhiều trong trường hợp này (thời gian gần như y hệt so với khi không bật).

Lý do có thể là do dữ liệu của chúng ta quá bé (ảnh 28x28) và model cũng quá nhỏ. Chi phí để pin bộ nhớ gần bằng với lợi ích nó mang lại. Tuy nhiên, với các mô hình lớn và dữ liệu đầu vào nặng (như ảnh độ phân giải cao), đây là một tham số bắt buộc phải bật để tối ưu tốc độ.

Refactor Dataset nhằm Giảm Chi phí CPU#

Ok, tiếp theo tôi sẽ nói đến một điểm tối ưu quan trọng khác. Sau khi dùng num_workers, tôi đã gỡ được bottleneck I/O, nhưng bây giờ bottleneck lại chuyển sang CPU processing.

Tôi nhận thấy đoạn code bên trong __getitem__ của class Dataset (v1) có một vấn đề rất nghiêm trọng:

# Bên trong __getitem__(self, idx)
row = self.df.iloc[idx]
pixels = row[self.feature_cols].to_numpy(dtype=np.float32)
image = torch.tensor(pixels.reshape(1, 28, 28)) / 255.0

Toàn bộ các thao tác: iloc (truy cập DataFrame), to_numpy, torch.tensor, reshape, và phép chia 255.0... toàn bộ đều được thực hiện khi một item bất kỳ được gọi.

Mặc dù self.df đã nằm trong RAM (không bị I/O bound), nhưng việc xử lý này được lặp đi lặp lại 60,000 lần mỗi epoch (thực hiện bởi 4 num_workers). Đây là một sự lãng phí tài nguyên CPU cực lớn.

Giải pháp: Nếu đã xác định load toàn bộ vào RAM, tại sao không thực hiện luôn các bước biến đổi trên một lần duy nhất tại thời điểm khởi tạo (__init__)?

Sau đó, __getitem__ chỉ có độ phức tạp thấp nhất là $O(1)$ (lấy phần tử đã xử lý từ một mảng/tensor).

Class Dataset (Version 2)#

Tôi đã implement lại như sau:

class MNISTDataset(Dataset):
    def __init__(self, parquet_path: str, num_classes: int = 10):
        df = pd.read_parquet(parquet_path)
        self.label_col = 'label' if 'label' in df.columns else None
        self.feature_cols = [c for c in df.columns if c != self.label_col]
        self.num_classes = num_classes
 
        # Xử lý toàn bộ features MỘT LẦN DUY NHẤT
        features_np = df[self.feature_cols].to_numpy(dtype=np.float32)
        # Reshape và chuyển thành Tensor MỘT LẦN
        self.images = torch.from_numpy(features_np.reshape(-1, 1, 28, 28)) / 255.0
 
        # Xử lý toàn bộ labels MỘT LẦN DUY NHẤT
        if self.label_col:
            labels_np = df[self.label_col].to_numpy(dtype=np.int64)
            self.labels = torch.from_numpy(labels_np).long()
        else:
            self.labels = torch.zeros(len(df), dtype=torch.long)
 
    def __len__(self):
        # Trả về độ dài đã được lưu
        return len(self.images)
    
    def __getitem__(self, idx):
        # Chỉ đơn giản là truy cập chỉ số (cực nhanh)
        image = self.images[idx]
        label = self.labels[idx]
        return {
            "image": image, 
            "label": label
        }

Result#

Bây giờ, hãy chạy lại với Dataset version 2, và giữ nguyên cấu hình DataLoader tối ưu (bs=512, num_workers=4, pin_memory=True).

Wow, và kết quả thật bất ngờ, giảm tới gần 9 lần so với bản tối ưu DataLoader trước đó.

Average Train Time per Epoch: 0.44s
Average Test Time per Epoch:  0.06s
Total Train Time: 2.21s
Total Test Time:  0.29s

Một epoch giờ chỉ mất 0.5 giây (0.44s train + 0.06s test) - một kết quả đáng kinh ngạc.

Tổng kết lại toàn bộ quá trình:

Kết luận#

Thật thú vị khi mà chúng ta tối ưu được tốc độ huấn luyện nhanh lên gấp khoảng 39 lần so với cách làm naive ban đầu. Nó hoàn toàn làm tôi quên mất là còn tập MNIST chưa đụng tới, nhưng mà thôi, tôi sẽ test hai tập này với các thuật toán vision mới ở các bài viết sau. Hy vọng là tôi lại khám phá ra một cách tăng tốc độ hơn nữa. ¹