#使用Pytorch进行简单的自定义图像分类 ~ONNX 推理

   图像分类是计算机视觉中的一项基本任务，涉及训练模型将图像分类为预定义类别。本文中，我们将探讨如何使用 PyTorch 构建一个简单的自定义对象分类模型，然后使用 ONNX 格式将其部署用于推理。

数据集准备

    在开始创建模型之前，准备一个标记数据集至关重要。收集要分类的不同对象类别的图像，并根据其类别将它们组织到单独的文件夹中。确保每个类别都有足够数量的图像，以避免过度拟合。

    准备如下树所示的数据集

- data
        - Fruits (dataset name)
            - train
                - class 1
                - class 2
                - ...
                - class n
            - val
                - class 1
                - class 2
                - ...
                - class n
            - test
                - class 1
                - class 2
                - ...
                - class n
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  • 10.
  • 11.
  • 12.
  • 13.
  • 14.
  • 15.
  • 16.
  • 17.

    我从 kaggle 获取了水果数据集，作为示例链接在此处：

https://www.kaggle.com/datasets/shreyapmaher/fruits-dataset-images
  • 1.

    更改 main.py 中“train_dir”和“val_dir”中的路径名

    如果需要，初始化数据加载器并添加增强功能。

构建模型

    首先导入必要的库，包括 PyTorch。定义自定义对象分类模型的架构，通常使用卷积神经网络 (CNN)。设计网络层，包括卷积层和池化层

import torch.nn as nn


class CustomConvNet(nn.Module):
    def __init__(self, num_classes):
        super(CustomConvNet, self).__init__()
        self.num_classes = num_classes
        self.layer1 = self.conv_module(3, 16)
        self.layer2 = self.conv_module(16, 32)
        self.layer3 = self.conv_module(32, 64)
        self.layer4 = self.conv_module(64, 128)
        self.layer5 = self.conv_module(128, 256)
        self.gap = self.global_avg_pool(256, self.num_classes)
    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = self.layer5(out)
        out = self.gap(out)
        out = out.view(-1, self.num_classes)
        return out
    def conv_module(self, in_num, out_num):
        return nn.Sequential(
            nn.Conv2d(in_num, out_num, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(out_num),
            nn.LeakyReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2))
    def global_avg_pool(self, in_num, out_num):
        return nn.Sequential(
            nn.Conv2d(in_num, out_num, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(out_num),
            nn.LeakyReLU(),
            nn.AdaptiveAvgPool2d((1, 1)))
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  • 10.
  • 11.
  • 12.
  • 13.
  • 14.
  • 15.
  • 16.
  • 17.
  • 18.
  • 19.
  • 20.
  • 21.
  • 22.
  • 23.
  • 24.
  • 25.
  • 26.
  • 27.
  • 28.
  • 29.
  • 30.
  • 31.
  • 32.
  • 33.
  • 34.

训练模型并进行评估

    使用随机梯度下降 (SGD) 或 Adam 优化器等技术实现训练循环，包括前向和后向传播、损失计算和梯度优化。通过跟踪损失和准确度等指标来监控训练过程。我们利用数据增强和正则化等技术来提高模型的泛化能力。

python main.py
  • 1.
for epoch in range(num_epochs):
    print("Epoch No -", epoch)


    model.train()
    running_loss = 0.0
    running_corrects = 0


    for inputs, labels in dataLoaders["train"]:
        # Feeding input and labels to device
        inputs = inputs.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)
        optimizer.zero_grad()
        with torch.set_grad_enabled(True):
            outputs = model(inputs)
            _, preds = torch.max(outputs,1)
            loss = criterion(outputs, labels)


            loss.backward()
            optimizer.step()
        running_loss += loss.item()* inputs.size(0)
        #calculate accuracy
        running_corrects += torch.sum(preds == labels.data)
    #scheduler step
    exp_lr_scheduler.step()
    # Calculate average loss and acc for a epoch
    epoch_loss = running_loss/len(train_data)
    epoch_acc = running_corrects.double()/len(train_data)


    print('Loss:{} , Acc{}'.format(epoch_loss, epoch_acc))
    # Saving model every five epoch
    if (epoch%5 == 0):
        save_model(model,epoch)
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  • 10.
  • 11.
  • 12.
  • 13.
  • 14.
  • 15.
  • 16.
  • 17.
  • 18.
  • 19.
  • 20.
  • 21.
  • 22.
  • 23.
  • 24.
  • 25.
  • 26.
  • 27.
  • 28.
  • 29.
  • 30.
  • 31.
  • 32.
  • 33.
  • 34.
  • 35.
  • 36.

    运行预测

    确保你是否已经更改了权重文件名、训练文件夹、predict.py 中的输入图像

python predict.py
  • 1.

导出为 ONNX 格式

    一旦您的模型经过训练并在验证集上表现良好，就可以将其导出为 ONNX 格式进行部署和推理。ONNX（开放神经网络交换）是一种开放标准格式，允许不同深度学习框架之间的互操作性。PyTorch 提供了将模型导出为 ONNX 的工具。

    确保你是否更改了export.py中的权重文件名

python export.py
  • 1.
# Now we will save this model.
import torch.onnx
torch.onnx.export(model,
                  img,
                  "./savedModels/custommodel.onnx",
                  export_params=True,
                  opset_version=10,
                  verbose=True,              # Print verbose output
                  input_names=['input'],     # Names for input tensor
                  output_names=['output'])
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  • 10.

使用 ONNX 进行推理

    加载已保存的 ONNX 模型并对新的未见过的图像进行推理。使用 ONNX 运行时库加载模型、提供输入数据并获得预测。测量推理时间并将其与 PyTorch 模型的推理时间进行比较，以评估通过 ONNX 优化实现的任何性能改进。

    确保你是否已经更改了权重文件名、训练文件夹、predict.py 中的输入图像

python onnx_inference.py
  • 1.
# Load the ONNX model
onnx_model = onnx.load("./savedModels/custommodel.onnx")


# Create an ONNX runtime session
ort_session = onnxruntime.InferenceSession("./savedModels/custommodel.onnx")


inputs = {"input": trans_img.numpy()}
outputs = ort_session.run(None, inputs)
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  • 10.


    在本教程中，我们探索了使用 PyTorch 构建简单自定义对象分类模型的过程。我们学习了如何训练模型、评估其性能并将其导出为 ONNX 格式以进行推理。通过利用 ONNX 运行时，我们演示了如何高效地对新图像进行推理。有了这些知识，您现在可以将自定义对象分类应用于各种实际应用程序并无缝部署模型。