当前位置：首页 > news >正文

10 卷积神经网络

news 2025/8/3 4:41:42

#----导包
import torch
from torch import nn
from torchvision import transforms
from torchvision import datasets
from torch.utils.data import DataLoader
import torch.nn.functional as F
import torch.optim as optim#----准备数据集
batch_size = 64
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307,),(0.3081,))])#----下载和加载train和test
trainset = datasets.MNIST(root='./data', train=True, transform=transform, download=False)
train_loader = DataLoader(trainset, batch_size=batch_size, shuffle=True)testset = datasets.MNIST(root='./data', train=False, transform=transform, download=False)
test_loader = DataLoader(trainset, batch_size=batch_size, shuffle=False)#-----搭建卷积神经网络
class Net(torch.nn.Module):def __init__(self):super().__init__()self.conv1 = nn.Conv2d(1, 10, 5)self.conv2 = nn.Conv2d(10, 20, 5)self.pool = nn.MaxPool2d(2, 2)self.fc1 = nn.Linear(320, 10)def forward(self, x):x = F.relu(self.pool(self.conv1(x)))x = F.relu(self.pool(self.conv2(x)))x = x.view(x.size(0), -1) #可以改成x = x.view(-1, 320)x = self.fc1(x)return xmodel = Net()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") #将模型迁移到GPU
model.to(device)#将模型迁移到GPUcriterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)def train(epoch):running_loss = 0.0for batch_idx, (inputs, targets) in enumerate(train_loader):inputs, targets = inputs.to(device), targets.to(device)optimizer.zero_grad()outputs = model(inputs)loss = criterion(outputs, targets)loss.backward()optimizer.step()running_loss += loss.item()if batch_idx % 300 == 299:print('[%d, %5d] loss: %.3f' %(epoch + 1, batch_idx + 1, running_loss / 300))running_loss = 0.0def test(epoch):correct = 0total = 0with torch.no_grad():for batch_idx, (inputs, targets) in enumerate(test_loader):inputs, targets = inputs.to(device), targets.to(device)outputs = model(inputs)_, predicted = torch.max(outputs.data, 1)total += targets.size(0)correct += predicted.eq(targets).sum().item()accuracy = 100 * correct / totalprint('Accuracy of the network on the 10000 test images: %f ' % accuracy)

训练和测试结果：