优化pytorch DataLoader提升数据加载速度

Posted 2023-03-18

参考技术A 因为pytorch数据加载速度太慢，影响训练速度，实训快速加载数据方式，提前获取要加载的数据，整体速度能快1/6.

操作步骤如下所示：

1、激活自己的torch虚拟环境:

source activate torch

2、安装prefetch_generator包

pip install prefetch_generator

3、定义DataLoaderX，继承torch原有的DataLoaderX的属性

class DataLoaderX(DataLoader):

    def __iter__(self):

        return BackgroundGenerator(super().__iter__())

4、使用的时候只需要把DataLoader变成DataLoaderX：

如下所示：

gen= DataLoaderX(train_dataset,shuffle=True,

batch_size=batch_size,

num_workers=4,pin_memory=True,

drop_last=False,

collate_fn=yolo_dataset_collate)

PyTorch - 优化模型参数

参考：OPTIMIZING MODEL PARAMETERS

梯度反向传播算法，更新参数

SGD -> Adam

dataset -> dataloader，train_dataloader训练，test_dataloader测试，迭代器

模型继承Module类，__init__(self)，定义层，Flatten()展平，和Sequential()，有序的容器

forward()函数，前向的计算，logits输出类别数，10个类别

Hyperparameters，超参数，不参与优化，影响模型的效果

Loss Function，损失函数、目标函数，分类函数CrossEntropyLoss，回归函数MSELoss

Optimizer，优化器，SGD，对参数进行更新，torch.optim.SGD()，更新模型参数，model.parameters()

优化前，调用optimizer.zero_grad()，计算梯度loss.backward()，optimizer.step()更新所有参数

optimizer.zero_grad()
loss.backward()
optimizer.step()

torch.no_grad()进行推理，计算正确率correct

数据集部分，编写dataset的自定义类，Transformer函数，Collection函数，dataset -> dataloader

model可以替换timm类，调用require_grad(False)，冻结参数

简单的分类或回归任务，seq2seq，训练传入真实值，测试是自回归任务

AI任务，数据集 + 模型 + 训练，预测时，是不需要优化参数的

PyTorch的torch.autograd()，自动微分

不使用zero_grad()，学习率需要降低，实际项目需要自定义数据集

序列建模，需要padding操作，有一些非法的值，Python类和类的继承，学习率 + 优化器

Embedding类，浮点类型，输入one-hot向量，权重就是embedding vector

源码：

• 使用timm模型替换基础模型
• 将灰度图像，转换为彩色图像

import timm
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor, Lambda

training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor()
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor()
)

train_dataloader = DataLoader(training_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
#         self.linear_relu_stack = nn.Sequential(
#             nn.Linear(28*28, 512),
#             nn.ReLU(),
#             nn.Linear(512, 512),
#             nn.ReLU(),
#             nn.Linear(512, 10),
#         )
        self.mobilenetv3 = timm.create_model('mobilenetv3_large_100', num_classes=10, pretrained=True)

    def forward(self, x):
        x = torch.cat([x, x, x], dim=1)  # 灰度图转换为彩色图
        logits = self.mobilenetv3(x)
        return logits

model = NeuralNetwork()

learning_rate = 1e-3
batch_size = 64
epochs = 5

# Initialize the loss function
loss_fn = nn.CrossEntropyLoss()

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: loss:>7f  [current:>5d/size:>5d]")


def test_loop(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0

    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \\n Accuracy: (100*correct):>0.1f%, Avg loss: test_loss:>8f \\n")

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

epochs = 10
for t in range(epochs):
    print(f"Epoch t+1\\n-------------------------------")
    train_loop(train_dataloader, model, loss_fn, optimizer)
    test_loop(test_dataloader, model, loss_fn)
print("Done!")