利用pytorch搭建lenet网络，并利用mnist数据集进行训练测试_怎么自己作一个lenet数据集资源-CSDN文库

175 浏览量 2021-01-06 16:11:10 上传评论 1 收藏 31KB PDF 举报

资源推荐

资源详情

资源评论

利用利用pytorch搭建搭建lenet网络，并利用网络，并利用mnist数据集进行训练测数据集进行训练测

试试

最近在学习pytorch，手工复现了LeNet网络，并附源码如下，欢迎大家留言交流

import torch

import torch.nn as nn

import torch.nn.functional as F

import torch.optim as optim

from torchvision import datasets, transforms

class LeNet(nn.Module):

def __init__(self):

super(LeNet,self).__init__()

self.conv1 = nn.Conv2d(1, 6, 5, 1,)

self.conv2 = nn.Conv2d(6, 16, 5, 1)

self.fc1 = nn.Linear(16*4*4, 120)

self.fc2 = nn.Linear(120, 84)

self.fc3 = nn.Linear(84, 10)

def forward(self, x):

#x:1*28*28

x = F.max_pool2d(self.conv1(x), 2, 2)

x = F.max_pool2d(self.conv2(x), 2, 2)

x = x.view(-1, 16*4*4)

x = self.fc1(x)

x = self.fc2(x)

x = self.fc3(x)

return F.log_softmax(x,dim=1)

def train(model, device, train_loader, optimizer, epoch):

model.train()

for idx, (data, target) in enumerate(train_loader):

data, target = data.to(device), target.to(device)

pred = model(data)

loss = F.nll_loss(pred,target)

optimizer.zero_grad()

loss.backward()

optimizer.step()

if idx % 100 == 0:

print("Train Epoch: {}, iterantion: {}, Loss: {}".format(epoch, idx, loss.item()))

def test(model, device, test_loader):

model.eval()

total_loss = 0.

correct = 0.

with torch.no_grad():

for idx, (data, target) in enumerate(test_loader):

data, target = data.to(device), target.to(device)

output = model(data)

total_loss += F.nll_loss(output, target, reduction="sum").item()

pred = output.argmax(dim=1)

correct +=pred.eq(target.view_as(pred)).sum().item()

total_loss /= len(test_loader.dataset)

acc = correct/len(test_loader.dataset)*100

print("Test loss: {}, Accuracy: {}".format(total_loss, acc))

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

batch_size = 32

train_dataloader = torch.utils.data.DataLoader(

datasets.MNIST("./MNIST_data", train=True, download=False,

transform = transforms.Compose([

transforms.ToTensor(),

本内容试读结束，登录后可阅读更多

下载后可阅读完整内容，剩余1页未读，立即下载

内容反馈

weixin_38617001

粉丝: 5
资源: 902

最新资源

资源上传下载、课程学习等过程中有任何疑问或建议，欢迎提出宝贵意见哦~我们会及时处理！点击此处反馈

feedback-tip