用Pytorch实现水果分类，训练+测试（内附数据集下载，完整代码）

一.数据集下载

链接: https://pan.baidu.com/s/1_7blbYJc0ouCGmqe8kBnTw 提取码: c6ex 复制这段内容后打开百度网盘手机App，操作更方便哦

二.训练模型

1.定义数据初始化

import torchvision.transforms as transforms
image_size=(224,224)
data_transforms=transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.Resize(image_size ),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),

])

一、 裁剪——Crop

1.随机裁剪：transforms.RandomCrop
2.中心裁剪：transforms.CenterCrop
3.随机长宽比裁剪 transforms.RandomResizedCrop
4.上下左右中心裁剪：transforms.FiveCrop
5.上下左右中心裁剪后翻转: transforms.TenCrop
二、翻转和旋转——Flip and Rotation

6.依概率p水平翻转transforms.RandomHorizontalFlip
7.依概率p垂直翻转transforms.RandomVerticalFlip
8.随机旋转：transforms.RandomRotation
三、图像变换

9.resize：transforms.Resize
10.标准化：transforms.Normalize
11.转为tensor：transforms.ToTensor
12.填充：transforms.Pad
13.修改亮度、对比度和饱和度：transforms.ColorJitter
14.转灰度图：transforms.Grayscale
15.线性变换：transforms.LinearTransformation()
16.仿射变换：transforms.RandomAffine
17.依概率p转为灰度图：transforms.RandomGrayscale
18.将数据转换为PILImage：transforms.ToPILImage
19.transforms.Lambda
四、对transforms操作，使数据增强更灵活

20.transforms.RandomChoice(transforms)
21.transforms.RandomApply(transforms, p=0.5)
22.transforms.RandomOrder

二.导入数据集

import torchvision.datasets as datasets
train_data=datasets .ImageFolder (root='fruits-360-original-size/fruits-360-original-size/Training',transform=data_transforms  )
print(train_data.classes)

输出为：

三.用DataLoader函数处理数据集

from torch.utils.data import DataLoader
batchsize=16
train_loader=DataLoader(train_data,batch_size=batchsize,shuffle=True,num_workers=1)

注意：
1.”batch_size= “一次训练所选取的样本数。
2.”shuffle=”输入一个布尔值来决定是否打乱数据
3.”num_workers= “是服务于多进程(multiprocessing)数据加载的, 用于设置有多少个子进程负责数据加载. num_workers并不是越大越好, 因为过多的子进程会占据 CPU 计算资源, 使得程序中其他在CPU上的计算部分变慢, 导致整体运行时间增加.

一般来说是通过逐步增加尝试来进行设置, 比如当GPU计算利用率已经很饱和时, 说明数据读取足够满足计算需求, 则不必再增加worker数量.

四.导入模型

import torchvision.models as models
AlexNet=models.AlexNet()
model=AlexNet

注意：
torchvision.models里面内置了许多预训练的模型，比如：
AlexNet
VGG
ResNet
SqueezeNet
DenseNet
Inception v3
其中ResNet18和AlexNet可以流畅运行。除了用Pytorch自带的模型也可以从本地导入模型，也可以自己定义模型。

五.权重初始化

import torch.nn.init as init
for name,module in model._modules.items() :
    if (name=='fc'):
        init.kaiming_uniform_(module.weight,a=0,mode='fan_in')

1.均匀分布 torch.nn.init.uniform_(tensor, a=0, b=1) 服从~U(a,b)
2.正太分布 torch.nn.init.normal_(tensor, mean=0, std=1) 服从~ N(mean,std)
3.初始化为常数 torch.nn.init.constant_(tensor, val) 初始化整个矩阵为常数val
4.Xavier
5.kaiming

六.定义优化器

import torch
optimizer=torch.optim .SGD(model.parameters(),lr=0.01 )
StepLR=torch.optim.lr_scheduler .StepLR (optimizer ,step_size= 3,gamma=1 )

注意：
StepLR（）作用是等间隔的调整学习率，调整倍数为gamma倍，调整间隔为step_size。

七.调用GPU

if torch.cuda.device_count()>1:
    model=nn.DataParallel(model)
print(device)

注意：
用GPU训练的模型，测试时只能用GPU测试

八.训练模型

import torch.nn.functional as F

def get_num_correct(outs, label):
    pass

model.to(device)

for epoch in range(10):
    total_loss=0
    print("epoch",epoch,":***************************")

    for batch in train_loader:
        images,labels=batch

        images=images.to(device)　
        labels=labels.to(device)

        outs=model(images)
        loss=F.cross_entropy(outs,labels)
        optimizer.zero_grad()　　
        loss.backward()
        optimizer .step()
        total_loss +=loss.item()
    print('loss',total_loss )

九.保存模型

torch.save(model,'FruitModelGPU.pth')

三.测试模型

1.定义数据初始化

import torchvision.transforms as transforms
image_size=(224,224)
data_transforms=transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.Resize(image_size ),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),

])

2.导入模型

import torch
model=torch.load('/home/ilab/FruitModelCpu4.pth',map_location='cpu')
model=model.module
model.eval()

注意：
因为使用GPU训练的模型，此时我们用CPU来进行测试，因此需要加上map_location=’cpu’，以及另起一行加上model=model.module

3.导入测试集(测试集为文件夹时)

import torchvision.datasets as datasets
test_data=datasets .ImageFolder (root='fruits-360-original-size/fruits-360-original-size/test',transform=data_transforms  )

4.进行测试

batch = next(iter(test_data))
images, labels = batch
out = model(images)
print(out)

注意：当测试集为打乱的文件时，前面的方法就行不通了，可以用下列方法：

1.定义数据初始化

import torchvision.transforms as transforms
image_size=(224,224)
data_transforms=transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.Resize(image_size ),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),

])

2.导入训练集

import torchvision.datasets as datasets
train_data=datasets .ImageFolder (root='/ilab/datasets/local/fruits/train',transform=data_transforms )

3.导入模型

import torch
model=torch.load('/home/ilab/FruitModelCpu4.pth',map_location='cpu')
model=model.module

4.进行测试

from PIL import Image
import os
i=1
sum=22688
root='/ilab/datasets/local/fruits/test'
file_list=os.listdir(root)
model.eval()
with open('/home/ilab/submission','w') as f:
        for filename in file_list:
            filepath = os.path.join(root,filename)
            img=Image.open(filepath)
            x=data_transforms(img).reshape((1,3,224,224))
            y= model(x)
            re=filename+' '+train_data.classes[int(y.argmax(dim=1))]+'\n'
            f.write(re)
            print(re)

5.FruitModelCpu4.pth下载地址：

链接: https://pan.baidu.com/s/17sa4Rda9jLhNRzKVkfsBMA 提取码: 2xv8 复制这段内容后打开百度网盘手机App，操作更方便哦

四.完整代码

"""""""""""""""""""""""定义数据初始化方法"""""""""""""""""""""""
import torch
import torch.nn
import torchvision
import torchvision.transforms as transforms
image_size=(224,224)
data_transforms=transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.Resize(image_size ),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),

])
"""""""""""""""""""""""导入数据集"""""""""""""""""""""""
import torchvision.datasets as datasets
train_data=datasets .ImageFolder (root='fruits-360-original-size/fruits-360-original-size/Training',transform=data_transforms  )
print(train_data.classes)
"""""""""""""""""""""制作DataLoader"""""""""""""""""""""""
from torch.utils.data import DataLoader
batchsize=500
train_loader=DataLoader(train_data,batch_size=batchsize,shuffle=True,num_workers=1)
""""""""""""""""""""""制作模型"""""""""""""""""""""""""""""
import torchvision.models as models
AlexNet=models.AlexNet()
model=AlexNet
"""""""""""""""""""""权重初始化"""""""""""""""""""""""
import torch.nn.init as init
for name,module in model._modules.items() :
    if (name=='fc'):
        init.kaiming_uniform_(module.weight,a=0,mode='fan_in')
"""""""""""""""""""""""定义优化器"""""""""""""""""""""
import torch
optimizer=torch.optim .SGD(model.parameters(),lr=0.01 )
StepLR=torch.optim.lr_scheduler .StepLR (optimizer ,step_size= 3,gamma=1 )
""""""""""""""""""""""调用GPU"""""""""""""""""""""""""""""

if torch.cuda.device_count()>1:
    model=nn.DataParallel(model)
print(device)
"""""""""""""""""""""""""""""""训练"""""""""""""""""""""""""""""""""
import torch.nn.functional as F

def get_num_correct(outs, label):
    pass

model.to(device)

for epoch in range(10):
    total_loss=0
    print("epoch",epoch,":***************************")

    for batch in train_loader:
        images,labels=batch

        images=images.to(device)
        labels=labels.to(device)

        outs=model(images)
        loss=F.cross_entropy(outs,labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer .step()
        total_loss +=loss.item()
    print('loss',total_loss )

""""""""""""""""""""""测试"""""""""""""""""""""""""
""""""""""""""""""""""权重初始化"""""""""""""""""""""""""
import torchvision.transforms as transforms
image_size=(224,224)
data_transforms=transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.Resize(image_size ),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
"""""""""""""""""""""""导入数据集"""""""""""""""""""""""
import os
import torchvision.datasets as datasets
train_data=datasets .ImageFolder (root='/ilab/datasets/local/fruits/train',transform=data_transforms )
"""""""""""""""""""""""导入模型"""""""""""""""""""""""
import torch
model=torch.load('/home/ilab/FruitModelCpu4.pth',map_location='cpu')
model=model.module
"""""""""""""""""""""""测试"""""""""""""""""""""""
 from PIL import Image
import os
i=1
sum=22688
root='/ilab/datasets/local/fruits/test'
file_list=os.listdir(root)
model.eval()
with open('/home/ilab/submission','w') as f:
        for filename in file_list:
            filepath = os.path.join(root,filename)
            img=Image.open(filepath)
            x=data_transforms(img).reshape((1,3,224,224))
            y= model(x)
            re=filename+' '+train_data.classes[int(y.argmax(dim=1))]+'\n'
            f.write(re)
            print(re)

])

Original: https://blog.csdn.net/qq_42548340/article/details/120825453
Author: Tomorrow；
Title: 用Pytorch实现水果分类，训练+测试（内附数据集下载，完整代码）