pytorch训练时一直打印迭代速度日志

原创

mob64ca12d0371b 2024-02-19 06:40:09 ©著作权

©著作权归作者所有：来自51CTO博客作者mob64ca12d0371b的原创作品，请联系作者获取转载授权，否则将追究法律责任

PyTorch训练时一直打印迭代速度日志

在使用PyTorch进行模型训练的过程中，我们经常会遇到需要查看每个迭代的速度日志的情况。这样可以帮助我们更好地监控模型训练的进度，以及及时发现问题。本文将介绍如何在PyTorch中实现一直打印迭代速度日志的功能，并提供代码示例。

实现方法

要实现一直打印迭代速度日志的功能，我们可以利用PyTorch中的DataLoader和torch.utils.data中的DataLoader类的batch_sampler属性。我们可以在DataLoader的迭代器中添加一个计时器，用于记录每个迭代的开始时间和结束时间，并计算迭代速度。然后在每个迭代结束时打印出速度日志。

代码示例

import time
import torch
from torch.utils.data import DataLoader

# 自定义DataLoader类，继承自torch.utils.data.DataLoader
class CustomDataLoader(DataLoader):
    def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None,
                 batch_sampler=None, num_workers=0, collate_fn=None,
                 pin_memory=False, drop_last=False, timeout=0, worker_init_fn=None):
        super(CustomDataLoader, self).__init__(dataset, batch_size, shuffle, sampler,
                                               batch_sampler, num_workers, collate_fn,
                                               pin_memory, drop_last, timeout, worker_init_fn)

    def __iter__(self):
        self.start_time = time.time()
        for batch in super(CustomDataLoader, self).__iter__():
            yield batch
            self.end_time = time.time()
            iter_time = self.end_time - self.start_time
            print(f"Iteration time: {iter_time} seconds")
            self.start_time = time.time()

# 使用自定义DataLoader
data_loader = CustomDataLoader(dataset, batch_size=32, shuffle=True)

for data in data_loader:
    # 模型训练过程
    pass

在上面的代码示例中，我们定义了一个CustomDataLoader类，继承自torch.utils.data.DataLoader。在__iter__方法中，我们记录了每个迭代的开始时间和结束时间，并计算了迭代时间。然后在每个迭代结束时打印出迭代时间。

示例应用

为了更好地演示迭代速度日志的功能，我们可以使用一个示例应用，比如训练一个简单的神经网络模型。以下是一个简单的示例代码：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as transforms

# 加载数据集
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = CustomDataLoader(train_dataset, batch_size=32, shuffle=True)

# 定义神经网络模型
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc = nn.Linear(28*28, 10)

    def forward(self, x):
        x = x.view(-1, 28*28)
        x = self.fc(x)
        return x

model = SimpleNN()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 开始训练
for data, target in train_loader:
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()

在上面的示例代码中，我们加载了MNIST数据集，并定义了一个简单的神经网络模型SimpleNN。然后使用自定义的CustomDataLoader进行训练。在训练过程中，会打印出每个迭代的速度日志。