pytorch实现自己的深度神经网络(公共数据集)

一、训练文件------train.py

注意:在运行此代码之前,需要配置好pytorch-GPU版本的环境,具体再次不谈。

python 复制代码
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

# 检查GPU是否可用
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Device:", device)

# 数据预处理的转换
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 加载CIFAR-10训练数据集
train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                               download=True, transform=transform)

# 创建数据加载器
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=8,
                                           shuffle=True, num_workers=0)

# 定义神经网络模型
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(128 * 32 * 32, 512)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = self.pool(torch.relu(self.conv3(x)))
        x = x.view(-1, 128 * 32 * 32)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 实例化模型,并将其移动到可用设备上
model = CNN().to(device)

# 定义损失函数
criterion = nn.CrossEntropyLoss()

# 定义优化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

if __name__ == '__main__':
    # 训练神经网络
    for epoch in range(5):
        running_loss = 0.0
        for i, data in enumerate(train_loader, 0):
            inputs, labels = data[0].to(device), data[1].to(device)

            # 梯度清零
            optimizer.zero_grad()

            # 正向传播
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            # 反向传播 + 优化
            loss.backward()
            optimizer.step()

            # 打印统计信息
            running_loss += loss.item()
            if i % 200 == 199:
                print('[%d, %5d] loss: %.3f' %
                      (epoch + 1, i + 1, running_loss / 200))
                running_loss = 0.0

    print('Finished Training')

    # 保存模型至文件
    torch.save(model.state_dict(), 'cifar10_cnn_model.pth')

二、测试文件------val.py

python 复制代码
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import cv2

# 检查GPU是否可用
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Device:", device)

# 数据预处理的转换
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

# 加载CIFAR-10测试数据集
test_dataset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                              download=True, transform=transform)

# 创建测试数据加载器
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=8,
                                             shuffle=False, num_workers=0)

# 加载模型并将其移动到可用设备上
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(128 * 32 * 32, 512)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = self.pool(torch.relu(self.conv3(x)))
        x = x.view(-1, 128 * 32 * 32)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x
# 显示函数
def imshow(img):
    img = img / 2 + 0.5
    npimg = img.numpy()
    # 坐标转换
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


model = CNN().to(device)
model.load_state_dict(torch.load('cifar10_cnn_model.pth'))
model.eval()


if __name__ == '__main__':
    # 在测试集上测试模型
    correct = 0
    total = 0
    with torch.no_grad():
        for data in test_loader:
            images, labels = data[0].to(device), data[1].to(device)
            outputs = model(images)
            # 预测值的最大值以及最大值的类别索引
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print('Accuracy on the test images: %d %%' % (100 * correct / total))

    # 显示测试集中的一些图片及其预测结果


    # 生成一个迭代器,从数据加载器中取出数据
    dataiter = iter(test_loader)
    # 从迭代器中获取下一个批次的数据
    images, labels = dataiter.next()
    # 将获取到的批次数据移动到device上,在这里也就是GPU上
    images, labels = images.to(device), labels.to(device)

    dip_flag = False
    if dip_flag == True:
        # -------------------------------------------
        # 可以选择 使用opencv显示
        # -------------------------------------------
        np_images = images.cpu().numpy()
        # 循环遍历并显示所有测试集图片
        for i in range(len(np_images)):
            # 从归一化中还原图像数据
            np_image = np.transpose(np_images[i], (1, 2, 0))   # 从CHW转换为HWC
            np_image = np_image * 0.5 + 0.5

            # 将图像数据从float类型转换为unit8类型
            np_image = (np_image * 255).astype(np.uint8)
            # 使用opencv显示图像
            cv2.imshow("Image {}".format(i+1), np_image)
            cv2.waitKey(0)
            # 等待用户按下任意键继续显示下一张图像
        cv2.destroyAllWindows()

    imshow(torchvision.utils.make_grid(images.cpu()))
    print('GroundTruth: ', ' '.join('%5s' % test_dataset.classes[labels[j]] for j in range(8)))
    outputs = model(images)
    _, predicted = torch.max(outputs, 1)
    print('Predicted: ', ' '.join('%5s' % test_dataset.classes[predicted[j]]
                                  for j in range(8)))

直接运行即可,亲测可以运行

相关推荐
G皮T3 小时前
【人工智能】ChatGPT、DeepSeek-R1、DeepSeek-V3 辨析
人工智能·chatgpt·llm·大语言模型·deepseek·deepseek-v3·deepseek-r1
九年义务漏网鲨鱼3 小时前
【大模型学习 | MINIGPT-4原理】
人工智能·深度学习·学习·语言模型·多模态
元宇宙时间3 小时前
Playfun即将开启大型Web3线上活动,打造沉浸式GameFi体验生态
人工智能·去中心化·区块链
开发者工具分享3 小时前
文本音频违规识别工具排行榜(12选)
人工智能·音视频
产品经理独孤虾4 小时前
人工智能大模型如何助力电商产品经理打造高效的商品工业属性画像
人工智能·机器学习·ai·大模型·产品经理·商品画像·商品工业属性
老任与码4 小时前
Spring AI Alibaba(1)——基本使用
java·人工智能·后端·springaialibaba
蹦蹦跳跳真可爱5894 小时前
Python----OpenCV(图像増强——高通滤波(索贝尔算子、沙尔算子、拉普拉斯算子),图像浮雕与特效处理)
人工智能·python·opencv·计算机视觉
雷羿 LexChien4 小时前
从 Prompt 管理到人格稳定:探索 Cursor AI 编辑器如何赋能 Prompt 工程与人格风格设计(上)
人工智能·python·llm·编辑器·prompt
两棵雪松5 小时前
如何通过向量化技术比较两段文本是否相似?
人工智能
heart000_15 小时前
128K 长文本处理实战:腾讯混元 + 云函数 SCF 构建 PDF 摘要生成器
人工智能·自然语言处理·pdf