加载自己的数据集训练，以VGG和LeNet为例，自己做的笔记，避免忘记

原创

已于 2023-09-15 14:45:12 修改 · 339 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#python #深度学习 #人工智能

于 2023-08-09 17:18:21 首次发布

本文介绍了如何使用PyTorch构建VGGNet模型对ImageNet数据集进行训练，包括数据预处理、模型定义、数据加载、训练过程和最终导出为ONNX格式以支持部署。

数据集为imagenet的格式，分为训练和验证两个文件夹，训练文件的子文件夹的名字代表具体的类别。

——ImageNet
  ——train
  	——cls1
  		——cls1_00.jpg
  		——cls1_01.jpg
  		...
  		——cls1_59.jpg
  	——cls2
  	——clsn
  ——test
  	——cls1
  		——cls1_60.jpg
  		——cls1_61.jpg
  		...
  		——cls1_100.jpg
  	——cls2
  	——clsn

import os
from torch.utils.data import Dataset, DataLoader
from PIL import Image
from torch import nn
import torch.optim as optim
from torchvision import transforms, models
import torch

class VGGNet(nn.Module):
    def __init__(self, num_classes=2):
        super(VGGNet, self).__init__()
        net = models.vgg16(pretrained=True)
        net.classifier = nn.Sequential()
        self.features = net
        self.classifier = nn.Sequential(
                nn.Linear(512 * 7 * 7, 512),
                nn.ReLU(True),
                nn.Dropout(),
                nn.Linear(512, 128),
                nn.ReLU(True),
                nn.Dropout(),
                nn.Linear(128, num_classes),
        )

    def forward(self, x):
        print("x.shape =" ,x.shape) #打印这行确定导出为onnx的输入的大小
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x) 
        return x

import torch.nn as nn

# 定义LeNet网络结构
# class LeNet(nn.Module):
#     def __init__(self, num_classes=2):
#         super(LeNet, self).__init__()
#         self.features = nn.Sequential(
#             nn.Conv2d(3, 6, kernel_size=5),
#             nn.