import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets,transforms
#datasets中提供了一些基本的数据集,transforms对图像进行基本的一些预处理
#导入一些库后,写一个简单的神经网络
class Net(nn.Module):
def __init__(self):
super(Net,self).__init__()
#conv2d参数分别是inpu_channel,output_c,kernel_size和stride
self.conv1 = nn.conv2d(1,20,5,1)
self.conv2 = nn.conv2d(20,50,5,1)
self.fc1 = nn.Linear(4*4*50,500)
self.fc2 = nn.Linear(500,10)
def forward(self,x):
x = F.relu(self.conv1(x)) #卷积后是20*24*24
#卷积后进行relu激活
x = F.max_pool2d(x,2,2) #池化后是20*12*12
#最大池化,池化层kernel大小和Stride均为2
x = F.relu(self.conv2(x)) #第二个卷积层
x = F.max_pooled(x,2,2) 4*4
#x.view在pytorch中相当于reshape
x = x.view(-1,4*4*50) #等价于x.view(x.shape[0],4*4*50)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return F.log_softmax(x,dim=1)
自动导入datasets中的一些数据集
#这样可以拿到该数据集
minst_data = datasets.MINST("./minst_data",train=true,download=True)
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) #batch_size*10
loss = F.nll_loss(pred,target)
#SGD
optimizer.zero(grad)
loss.backward()
optimizer.step()
#TEST
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)
#reduction = sum,对loss进行累加
total_loss += F.nll_loss(output,target,reduction = sum)
pred = oputput.argmax(dim=1) #batch_size*10
#拿出minst_data中的所有数据的第0维,将其转化为numpy
data = [d[0].data.cpu().numpy() for d in minst_data]
#准备数据到GPU或CPU上
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
batch_size = 32
#将数据转化为dataloader,就可以迭代
train_dataloader = torch.utils.data.DataLoader(
datasets.MINST("./minst_data",train = True, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
#预处理,标准化
transforms.Normalize((0.1307,),(0.3081,))
])),
#transforms对图像进行基本的一些预处理。 Compose就是就是将一些operation进行拼接,然后用.ToTensor()将该数据集全部转化为Tensor
batch_size = batch_size,shufftle =True,
num_workers=1, #线程数
pin_memory=True #据说pin_memory可加速过程
)
test_dataloader = torch.utils.data.DataLoader(
dataset.MINST("./minst_data",train=False,download=True,
transform = transforms.Compose([
transforms.ToTensor(),
#预处理,标准化
transforms.Normalize((0.1307,),(0.3081,))
]))
batch_size = batch_size,shufftle =True, num_workers=1,
pin_memory=True #据说pin_memory可加速过程
)
learning_rate = 0.01
momentum =0.5
model = Net().to(device)
optimizer = torch.optim.SGD(model.parameters(), lr = learning_rate,momentum = momentum)
num_epochs = 5
for epoch in range(num_epochs):
train(model,device,train_dataloader,optimizer,epoch)
test(model,device,test_dataloader)
#保存一下模型
torch.save(model.state_dict(),"minst_cnn.pt")