图像分类入门-CSDN博客

本文链接：https://blog.csdn.net/qq_36784503/article/details/135285892
１，简单的分类
例子
import torch
import torchvision
import torchvision.transforms as transforms
from tqdm import tqdm
import cv2 as cv
import numpy as np
import time
#--
import torch.nn as nn
import torch.nn.functional as F

import torch.optim as optim

#-------------------------------------------------------------data
# 定义归一化方法
# transform = transforms.Compose(
# 	# 首先装换数据为tensor张量
#     [transforms.ToTensor(),
#     # 对数据进行正态分布归一化，RGB三个通道每个通道均值为0.5，标准差为0.5
#      transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
#     ]
#     )
transform = transforms.Compose(
	# 首先装换数据为tensor张量
    [transforms.ToTensor()
    ]
    )

# 加载训练数据集
# root为数据存放的目录，train=Ture表示训练集，download=True表示要下载，transform为之前定义的归一化方法
# trainset = torchvision.datasets.CIFAR10(root='./data', train=True, 
#                                        download=True, transform=transform)
privateset = torchvision.datasets.ImageFolder(root="data/tire/train/", transform=transform)
 
# 定义数据集的加载方法， trainset为训练集， batch_size为单词训练的样本数，shuffle=True表示随机抽取样本，num_workers表示加载的线程数
trainloader = torch.utils.data.DataLoader(privateset, batch_size=64,
                                         shuffle=True, num_workers=1)


#-----------------------------------------model--

 
# 定义神经网络类
class Net(nn.Module):
	#定义神经网络结构, 输入数据 1x32x32
    def __init__(self): 
        super(Net, self).__init__()

        # 第一层（卷积层）
        # 输入频道3， 输出频道6， 卷积3x3
        # self.conv1 = nn.Conv2d(3,6,3) 
        self.conv1 = nn.Conv2d(
            in_channels=3,
            out_channels=45,
            kernel_size=(3, 3),
            stride=(2, 2),
            padding=(1, 1),
        ) # w=256
        # 第二层（卷积层）
        # 输入频道6， 输出频道16， 卷积3x3

        # 第三层（全连接层）
        self.conv2 = nn.Conv2d(
            in_channels=45,
            out_channels=30,
            kernel_size=(3, 3),
            stride=(2, 2),
            padding=(1, 1),
        ) # w =128
        self.conv3 = nn.Conv2d(
            in_channels=30,
            out_channels=60,
            kernel_size=(3, 3),
            stride=(2, 2),
            padding=(1, 1),
        ) # w=64
        self.conv4 = nn.Conv2d(
            in_channels=60,
            out_channels=16,
            kernel_size=(3, 3),
            stride=(2, 2),
            padding=(1, 1),
        ) # w=32 
        self.conv5 = nn.Conv2d(
            in_channels=16,
            out_channels=16,
            kernel_size=(3, 3),
            stride=(2, 2),
            padding=(1, 1),
        ) # w=16

        # 输入维度16x28x28=12544，输出维度 512
        self.fc1 = nn.Linear(16*16*16, 512) 
        # 第四层（全连接层）
        # 输入维度512， 输出维度64
        self.fc2 = nn.Linear(512, 64) 
        # 第五层（全连接层）
        # 输入维度64， 输出维度10
        self.fc3 = nn.Linear(64, 2) 
        
    # 定义数据流向
    def forward(self, x): 
    	# 数据先经过第一层卷积层
        x = self.conv1(x)
        # 经过激活函数
        x = F.relu(x)
        
        # 数据经过第二层卷积层
        x = self.conv2(x)
        # 经过激活函数
        x = F.relu(x)

        x = self.conv3(x)
        # 经过激活函数
        x = F.relu(x)

        x = self.conv4(x)
        # 经过激活函数
        x = F.relu(x)

        x = self.conv5(x)
        # 经过激活函数
        x = F.relu(x)
        
        # 调整数据维度，‘-1’表示自动计算维度
        x = x.view(-1, 16*16*16)
        # 数据经过第三层全连接层
        x = self.fc1(x)
        # 数据经过激活函数
        x = F.relu(x)
        
         # 数据经过第四层全连接层
        x = self.fc2(x)
        # 数据经过激活函数
        x = F.relu(x)
        
        # 数据经过第五层全连接层，输出结果
        x = self.fc3(x)
        
        return x
net = Net()
print(net)
net.load_state_dict(torch.load('./model_ok.pt'))
# 将模型及其子模型的参数全部放到 GPU 上
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device="cuda"
net_cuda=net.to(device)
# print(net)

 
# 导入torch中优化器相关的包

 
# 定义损失函数为交叉熵函数
criterion = nn.CrossEntropyLoss()
 
# 优化器函数为随机梯度下降， 学习率为0。0001
running_loss=0
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
start_time = time.time()
for epoch in range(500):
    #训练
    for i, data in enumerate(trainloader):
    	# 获得数据与标签
        images, labels = data
        # image_sub = np.array(images[0])
        # image_sub = np.transpose(image_sub, (1, 2, 0))
        # label_sub = np.array(labels[0])
        # label_string=np.array2string(label_sub)
        # print("labels",label_sub.shape)
        # print(label_sub)
        # print(image_sub.shape)
        # cv.imwrite("0.jpg",image_sub)
##------------------------put lab
##----------------------------------
        # cv.imshow("0",image_sub)
        # cv.waitKey(0) 
        images, labels = images.to(device), labels.to(device)
        # print("images.shape",images.shape)       
        # 得到网络的输出
        # net=net.to(device)
        outputs = net_cuda(images)
        # 计算损失
        # print("outputs:",outputs)
        # print("labels:",labels)
        loss = criterion(outputs, labels) 
        # 清零梯度
        optimizer.zero_grad()
        # 反向传播
        loss.backward()
        # 更新权重
        optimizer.step()
        # 计算总的损失
        running_loss += loss.item()
        #每1000 mini batch 测试一次
        if(i%150 == 0):
            print('Epoch: %d,Step: %d, Loss: %.3f'%(epoch,i,loss.item()) ) 
        end_time = time.time()
        
        
print ("time elapse:",end_time - start_time,"s")
torch.save(net.state_dict(), './model_2.pt')