from torch.utils.data import Dataset
from PIL import Image
import os
class myData(Dataset):
    def __init__(self, strRootDir, strLabelDir):
        self.mstrRootDir = strRootDir
        self.mstrLabelDir = strLabelDir
        self.mstrPath = os.path.join(strRootDir, strLabelDir)
        self.mstrImagePath = os.listdir(self.mstrPath)
    def __getitem__(self, item):
        strImageName = self.mstrImagePath[item]
        strImageItemPath = os.path.join(self.mstrPath, strImageName)
        zImage = Image.open(strImageItemPath)
        strLabel = self.mstrLabelDir
        return zImage, strLabel
    def __len__(self):
        return len(self.mstrImagePath)
strRootDir = "dateset/train"
strLabelAnts = "ants"
strLabelBees = "bees"
zAntsData = myData(strRootDir, strLabelAnts)
zBeesData = myData(strRootDir, strLabelBees)
zTrainData = zAntsData + zBeesData
print(len(zAntsData), len(zBeesData), len(zTrainData))

from torch.utils.tensorboard import SummaryWriter
import numpy
from PIL import Image
wirter = SummaryWriter("logs")
PILImagePath = "dateset/train/ants/0013035.jpg"
PILImage = Image.open(PILImagePath)
ImageArray = numpy.array(PILImage)
wirter.add_image("ants", ImageArray, 1, dataformats='HWC')
for i in range(100):
    wirter.add_scalar("y = 2x", 2 * i, i)
wirter.close()

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.model = nn.Sequential(
            # 卷积层1
            nn.Conv2d(3, 32, 5, 1, 2, ),
            # 池化层1
            nn.MaxPool2d(2),
            # 卷积层2
            nn.Conv2d(32, 32, 5, 1, 2),
            # 池化层2
            nn.MaxPool2d(2),
            # 卷积层3
            nn.Conv2d(32, 64, 5, 1, 2),
            # 池化层3
            nn.MaxPool2d(2),
            # 展平层
            nn.Flatten(),
            # 全连接层1
            nn.Linear(64 * 4 * 4, 64),
            # 全连接层2
            nn.Linear(64, 10)
        )
    def forward(self, x):
        """
        神经元向前传播函数
        :param x: 输入的参数
        :return: 输出的参数
        """
        x = self.model(x)
        return x

#!/usr/bin/env pytorch
# -*- coding: UTF-8 -*-
"""
@Project     :llearn_pytorch
@File        :model.py
@IDE         :PyCharm
@Author      :张世航
@Date        :2023/2/24 11:30
@Description :一个深度学习演示样例
"""
import torchvision.datasets
from torch.optim import lr_scheduler, SGD
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from model import *
import time
# 设置设备
zDevice = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 加载数据集
zTrainData = torchvision.datasets.CIFAR10("TrainData", train=True, transform=torchvision.transforms.ToTensor(),
                                          download=True)
zTestData = torchvision.datasets.CIFAR10("TestData", train=False, transform=torchvision.transforms.ToTensor(),
                                         download=True)
# 添加记录者
writer = SummaryWriter("logs")
# 计算数据集长度
iTrainDataLength = len(zTrainData)
iTestDataLength = len(zTestData)
print("the length of train data :{}".format(iTrainDataLength))
print("the length of test data:{}".format(iTestDataLength))
# 利用dataloader加载数据集
zTraindataLoader = DataLoader(zTrainData, batch_size=64)
zTestDataLoader = DataLoader(zTestData, batch_size=64)
# 创建网络模型
myModel = MyModel()
myModel = myModel.to(zDevice)
# 创建损失函数
myLossFunction = nn.CrossEntropyLoss()
myLossFunction = myLossFunction.to(zDevice)
# 优化器
dLearnRate = 1e-2
myOptimizer = SGD(myModel.parameters(), lr=dLearnRate)
# 设置学习率衰减函数
MyScheduler = lr_scheduler.StepLR(myOptimizer, 50, gamma=0.5)
# 设置训练网络的一些参数
# 训练的总次数
iTotalTrainStep = 0
# 测试的总次数
iTotalTestStep = 0
# 训练的轮数
iEpoch = 300
# 记录开始时间
fStartTime = time.time()
for i in range(iEpoch):
    print("----第{}训练开始！！!----".format(i))
    myModel.train()
    for data in zTraindataLoader:
        images, targets = data
        images = images.to(zDevice)
        targets = targets.to(zDevice)
        outputs = myModel(images)
        loss = myLossFunction(outputs, targets)
        # 优化器优化模型
        myOptimizer.zero_grad()
        loss.backward()
        myOptimizer.step()
        iTotalTrainStep = iTotalTrainStep + 1
        if iTotalTrainStep % 100 == 0:
            fEndTime = time.time()
            print("第{}次模型训练loss是{}".format(iTotalTrainStep, loss.item()))
            writer.add_scalar("train_loss", loss.item(), iTotalTrainStep)
            print("训练耗时{}".format(fEndTime - fStartTime))
    MyScheduler.step()
    print("----调整学习率为{}----".format(myOptimizer.state_dict()['param_groups'][0]['lr']))
    writer.add_scalar("train_lr", myOptimizer.state_dict()['param_groups'][0]['lr'], i)
    myModel.eval()
    iTotalLoss = 0
    iTotalAccuracy = 0
    print("----第{}测试开始！！!----".format(i))
    with torch.no_grad():
        for data in zTestDataLoader:
            images, targets = data
            images = images.to(zDevice)
            targets = targets.to(zDevice)
            outputs = myModel(images)
            loss = myLossFunction(outputs, targets)
            iTotalLoss = iTotalLoss + loss.item()
            accuracy = (outputs.argmax(1) == targets).sum()
            iTotalAccuracy = accuracy + iTotalAccuracy
    print("第{}次模型测试loss是{}".format(i, iTotalLoss))
    print("第{}次模型测试正确率是{}".format(i, iTotalAccuracy / iTestDataLength))
    writer.add_scalar("test_loss", iTotalLoss, iTotalTestStep)
    writer.add_scalar("test_accuracy", iTotalAccuracy / iTestDataLength, iTotalTestStep)
    iTotalTestStep = iTotalTestStep + 1
    # 保存每一次训练的模型
    torch.save(myModel.state_dict(), "model/model_{}.path".format(i))
    print("----模型已经保存！！！----")
writer.close()

#!/usr/bin/env pytorch
# -*- coding: UTF-8 -*-
"""
@Project     :llearn_pytorch 
@File        :testmodel.py
@IDE         :PyCharm 
@Author      :张世航
@Date        :2023/2/27 8:48 
@Description :一个简易的验证训练好的模型的程序
"""
import os
import torch
from PIL import Image
from torchvision import transforms
from model import MyModel
class myImage:
    def __init__(self, strRootDir):
        self.mstrRootDir = strRootDir
        self.mstrImagePath = os.listdir(self.mstrRootDir)
    def __getitem__(self, item):
        strImageName = self.mstrImagePath[item]
        strImageItemPath = os.path.join(self.mstrRootDir, strImageName)
        zImage = Image.open(strImageItemPath)
        strLabel = strImageName
        return zImage, strLabel
    def __len__(self):
        return len(self.mstrImagePath)
funTransform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor()
])
model = MyModel()
model.load_state_dict(torch.load("model/model_37.path"))
model.eval()
strTestDir = "testimage"
zData = myImage(strTestDir)
image_type = ("airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck", "nolen")
with torch.no_grad():
    for data in zData:
        image, label = data
        image = image.convert('RGB')
        image = funTransform(image)
        image = torch.reshape(image, (1, 3, 32, 32))
        output = model(image)
        iResult = output.argmax(1)
        print("图片类型|模型识别出类型:{}|{}".format(label, image_type[iResult]))