代码详解
batch_size = x.size(0)
这一行代码是用来获取当前批次(batch)的数据量。在 PyTorch 中,张量 x
的维度通常按照 [batch_size, channel, height, width]
的顺序排列,其中 batch_size
表示一批次中包含的样本数量。
x = x.view(batch_size, -1)
这一行中用于将卷积层的输出展平成一个一维向量,其中 -1
表示自动计算除 batch_size
外的其他维度,即这里的 320。这样就可以将展平后的向量输入到全连接层中了。
import torch
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
# prepare dataset
# 把原始图像转换为图像张量,C*W*H 然后对数据进行归一化,类比正态分布转化为标准形式
batch_size = 64
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,),(0.3081,))])
train_dataset = datasets.MNIST(root='../dataset/mnist/', train=True, download=True,transform=transform)
train_loader = DataLoader(train_dataset, shuffle=True,batch_size=batch_size)
test_dataset = datasets.MNIST(root='../dataset/mnist/', train=False, download=True,transform=transform)
test_loader = DataLoader(test_dataset, shuffle=False, batch_size=batch_size)
# design model using class
class Net(torch.nn.Module):
def __init__(self):
super(Net,self).__init__()
self.conv1 = torch.nn.Conv2d(1, 10, kernel_size=5) # 卷积核的大小是5*5
self.conv2 = torch.nn.Conv2d(10, 20, kernel_size=5)
self.pooling = torch.nn.MaxPool2d(2) # 提供2维的最大池化操作
# 定义全连接层
self.fc = torch.nn.Linear(320,10) # 输入大小为320(假设是展平后的向量长度),输出大小为10(分类任务中的类别数量)
def forward(self,x):
# flatten
batch_size = x.size(0) #获取当前批次的数据量
x = F.relu(self.pooling(self.conv1(x)))
x = F.relu(self.pooling(self.conv2(x)))
x = x.view(batch_size, -1)
x = self.fc(x)
return x
model = Net()
# 调用GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
def train(epoch):
running_loss = 0.0
for batch_idx, data in enumerate(train_loader, 0):
inputs, target = data
inputs, target = inputs.to(device), target.to(device) #GPU
outputs = model(inputs)
loss = criterion(outputs, target)
optimizer.zero_grad()
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))
def Test():
correct = 0
total = 0
with torch.no_grad(): # 禁用梯度计算。在测试集上不需要反向传播,禁用梯度计算减少内存消耗
for data in test_loader:
images, labels =data
images, labels = images.to(device), labels.to(device) # GPU
outputs = model(images)
_, predicted = torch.max(outputs.data, dim=1) # _说明不关注类别本身。只关注索引标签,在这里把标签预测的标签记为:predicted
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('accuracy on test set: %d %%' % (100*correct/total))
if __name__ == '__main__':
for epoch in range(10):
train(epoch)
Test()