Pytorch – Một số tips hay, tối ưu cho quá trình huấn luyện model của bạn

Xin chào các bạn, chúng ta đã lâu không gặp nhau trên viblo. Bây giờ mình quay lại để chia sẻ một số mẹo hay về Pytorch để tối ưu quá trình huấn luyện model của bạn. Mình hi vọng rằng những mẹo này sẽ giúp bạn đạt được kết quả tốt hơn và giảm thiểu các vấn đề như “out of memory”.

Chia Sẻ Về Tips và Trick

Trước khi bắt đầu, hãy cùng mình nhìn vào một bài toán kinh điển mà chúng ta đã từng làm ít nhất một lần: “Nhận Dạng Chữ Số Viết Tay với Bộ Dữ Liệu MNIST”. Đây là một bài toán thông thường chúng ta có thể sử dụng để thực hiện các thử nghiệm và tối ưu hóa trong Pytorch.

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 64)
        self.fc2 = nn.Linear(64, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output

Tiếp theo, chúng ta xây dựng hàm training cho model này:

def train_normal(model, device, train_loader, val_loader, optimizer, epoch):
    train_loss = []
    train_acc = []
    model.train()

    for (data, target) in train_loader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        acc = accuracy(output, target)
        train_loss.append(loss)
        train_acc.append(acc)
        loss.backward()
        optimizer.step()

    val_loss = []
    val_acc = []
    model.eval()

    for (data, target) in val_loader:
        data, target = data.to(device), target.to(device)
        output = model(data)
        loss = F.nll_loss(output, target)
        acc = accuracy(output, target)
        val_loss.append(loss)
        val_acc.append(acc)

    result = dict()
    result['train_loss'] = torch.stack(train_loss).mean()
    result['train_acc'] = torch.stack(train_acc).mean()
    result['val_loss'] = torch.stack(val_loss).mean()
    result['val_acc'] = torch.stack(val_acc).mean()
    print("Epoch [{}], train_loss: {:.4f}, train_acc: {:.4f}, val_loss: {:.4f}, val_acc: {:.4f}".format(
        epoch, result['train_loss'], result['train_acc'], result['val_loss'], result['val_acc']))

    return result

Đoạn mã trên sử dụng mô hình CNN cơ bản với hai lớp tích chập và hai lớp kết nối đầy đủ. Chúng ta sử dụng optimizer là Adadelta và hàm mất mát là Negative Log Likelihood Loss (NLLLoss).

Tuy model này có ít tham số (khoảng 609,354 tham số), và lưu trữ chỉ khoảng 2MB, nhưng vẫn có thể gặp lỗi “out of memory” đặc biệt với những máy tính có phần cứng yếu. Vậy chúng ta sẽ cùng nhau tìm hiểu một số tips và trick để tối ưu quá trình này.

DataLoader và Các Tham Số Quan Trọng

DataLoader là một module được Pytorch xây dựng nhằm hỗ trợ load và xử lí dữ liệu theo từng batch. Có một số tham số quan trọng mà chúng ta cần chú ý khi sử dụng DataLoader.

torch.utils.data.DataLoader(
    dataset: Dataset[T_co],
    batch_size: Optional[int] = 1,
    shuffle: bool = False,
    sampler: Optional[Sampler[int]] = None,
    batch_sampler: Optional[Sampler[Sequence[int]]] = None,
    num_workers: int = 0,
    collate_fn: Optional[_collate_fn_t] = None,
    pin_memory: bool = False,
    drop_last: bool = False,
    timeout: float = 0,
    worker_init_fn: Optional[_worker_init_fn_t] = None,
    multiprocessing_context=None,
    generator=None,
    *,
    prefetch_factor: int = 2,
    persistent_workers: bool = False
)

Có hai tips quan trọng mà chúng ta nên biết:

  1. Để tận dụng sức mạnh của máy tính, chúng ta có thể sử dụng nhiều workers (num_workers) để xử lí song song nhiều dữ liệu cùng một lúc. Điều này giúp giảm thời gian xử lí xuống đáng kể.

  2. Batch size ảnh hưởng đến chất lượng học của mô hình. Batch size càng lớn, mô hình sẽ học tốt hơn. Vì vậy, một tips quan trọng là tăng kích thước batch size lên tối đa có thể, sau đó giảm dần nếu gặp vấn đề về RAM hoặc tìm các phương pháp tối ưu bộ nhớ khác.

Với đoạn mã train_normal ở trên, chúng ta có thể thử nghiệm với các giá trị khác nhau cho batch size và num_workers, nhưng vẫn có thể gặp lỗi “out of memory”.

Sử Dụng 16-Bit Precision

Một trong những vấn đề tiêu biểu khiến chúng ta gặp lỗi “out of memory” là việc sử dụng kiểu 32-bit precision cho các tensor, dẫn đến việc máy tính phải sử dụng nhiều RAM hơn để lưu trữ. Tuy nhiên, hiện nay, Pytorch hỗ trợ kiểu 16-bit precision, giúp giảm đi 1 nửa dung lượng lưu trữ và tăng tốc độ tính toán mà vẫn đảm bảo hiệu suất và độ chính xác của mô hình.

Chúng ta có thể sử dụng thư viện torch.cuda.amp đã được tích hợp trong Pytorch từ phiên bản 1.6 trở lên để sử dụng 16-bit precision. Việc sử dụng torch.cuda.amp khá đơn giản với hai module chính: auto_cast và GradScaler.

scaler = GradScaler()

optimizer.zero_grad()
for i, (data, target) in enumerate(train_loader):
    with autocast():
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        acc = accuracy(output, target)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

Tích Lũy Đạo Hàm (Accumulated Gradients)

Việc tính toán và cập nhật gradients của mô hình thông qua quá trình forward và backward có thể gây áp lực lên bộ nhớ. Tuy nhiên, Pytorch cho phép chúng ta tích lũy gradients để giảm thiểu vấn đề này.

Để tích lũy gradients, thay vì cập nhật trọng số sau mỗi batch, chúng ta có thể đợi load n batch rồi mới thực hiện cập nhật một lần.

scaler = GradScaler()

optimizer.zero_grad()
for i, (data, target) in enumerate(train_loader):
    with autocast():
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        acc = accuracy(output, target)
        scaler.scale(loss).backward()

        # backward chỉ được thực hiện sau n batch
        if (i+1) % n == 0:
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

Lưu Trữ Loss và Accuracy

Việc lưu trữ loss và accuracy là điều cần thiết khi bạn muốn vẽ biểu đồ để quan sát quá trình training của mô hình. Tuy nhiên, việc lưu nguyên những tensor này có thể gây lỗi “out of memory”.

Để khắc phục vấn đề này, chúng ta chỉ cần chuyển các tensor thành float trước khi lưu trữ bằng cách thêm “.item()” vào cuối.

result = dict()
result['train_loss'] = torch.stack(train_loss).mean().item()
result['train_acc'] = torch.stack(train_acc).mean().item()
result['val_loss'] = torch.stack(val_loss).mean().item()
result['val_acc'] = torch.stack(val_acc).mean().item()

Đổi Tên Layer

Khi bạn muốn thay đổi tên của một layer, chúng ta cần cập nhật lại trọng số tương ứng của nó. Để thực hiện điều này mà không cần train lại toàn bộ model, chúng ta có thể thực hiện các bước sau:

  1. Khai báo model cũ và load weight cũ.
old_model = Net()
old_model.load_state_dict('old_weights.pt')
  1. Khởi tạo một state dict mới từ state dict cũ và cập nhật tên layer tương ứng.
old_state_dict = old_model.state_dict()
new_state_dict = copy.deepcopy(old_state_dict)

for key in old_state_dict:
    if 'conv' in key:
        new_state_dict[key.replace('conv', 'embbed')] = new_state_dict.pop(key)
  1. Đổi tên layer cho model mới.
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.embded_1 = nn.Conv2d(1, 32, 3, 1)
        self.embded_2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 64)
        self.fc2 = nn.Linear(64, 10)

    def forward(self, x):
        x = self.embded_1(x)
        x = F.relu(x)
        x = self.embded_2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output
  1. Khai báo model mới và load weight mới.
new_model = Net()
new_model.load_state_dict(new_state_dict)
torch.save(new_model.state_dict(), 'new_weights.pt')

Đó là những mẹo mà mình muốn chia sẻ trong bài viết này. Hy vọng rằng những mẹo và trick này sẽ giúp bạn trong quá trình huấn luyện mô hình. Nếu bạn thấy bài viết hay và hữu ích, hãy ủng hộ bằng cách upvote, clip và share. Nếu bạn có những mẹo khác, hãy comment dưới bài viết để chúng ta cùng thảo luận.

FEATURED TOPIC