import torch
|
import torch.nn as nn
|
import torchvision.datasets as dataset
|
import torchvision.transforms as transforms
|
import torch.utils.data as data_utils
|
import matplotlib.pyplot as plt
|
|
print(torch.__version__)
|
print(torch.cuda.is_available())
|
|
# 获取数据集
|
train_data = dataset.MNIST(root="D:\\ml",
|
train=True,
|
transform=transforms.ToTensor(),
|
download=True
|
)
|
test_data = dataset.MNIST(root="D:\\ml",
|
train=False,
|
transform=transforms.ToTensor(),
|
download=False
|
)
|
train_loader = data_utils.DataLoader(dataset=train_data, batch_size=100, shuffle=True)
|
test_loader = data_utils.DataLoader(dataset=test_data, batch_size=100, shuffle=True)
|
|
|
|
cnn = CNN()
|
# cnn = cnn.cuda()
|
|
# 损失函数
|
los = torch.nn.CrossEntropyLoss()
|
|
# 优化函数
|
optime = torch.optim.Adam(cnn.parameters(), lr=0.01)
|
|
train_targets = train_data.targets
|
plt.figure(figsize=(9, 9))
|
for i in range(9):
|
plt.subplot(3, 3, i + 1)
|
plt.title(train_targets[i].numpy())
|
plt.axis('off')
|
plt.imshow(train_data.data[i], cmap='gray')
|
plt.show()
|
|
# 训练模型
|
for epo in range(10):
|
for i, (images, lab) in enumerate(train_loader):
|
# images = images.cuda()
|
# lab = lab.cuda()
|
out = cnn(images)
|
loss = los(out, lab)
|
optime.zero_grad()
|
loss.backward()
|
optime.step()
|
print("epo:{},i:{},loss:{}".format(epo + 1, i, loss))
|
|
# 测试模型
|
loss_test = 0
|
accuracy = 0
|
with torch.no_grad():
|
for j, (images_test, lab_test) in enumerate(test_loader):
|
# images_test = images_test.cuda()
|
# lab_test = lab_test.cuda()
|
out1 = cnn(images_test)
|
loss_test += los(out1, lab_test)
|
loss_test = loss_test / (len(test_data) // 100)
|
_, p = out1.max(1)
|
accuracy += (p == lab_test).sum().item()
|
accuracy = accuracy / len(test_data)
|
print("loss_test:{},accuracy:{}".format(loss_test, accuracy))
|
