TensorFlow与Pytorch的转换——2手写数字识别
数据处理
import tensorflow as tf
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
# 将像素值缩放到0到1之间
x_train, x_test = x_train / 255.0, x_test / 255.0
# 将标签转换为one-hot编码
y_train = tf.keras.utils.to_categorical(y_train, num_classes=10)
y_test = tf.keras.utils.to_categorical(y_test, num_classes=10)
构建模型
model = tf.keras.models.Sequential([
tf.keras.layers.Flatten(input_shape=(28, 28)),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dense(64, activation='relu'),
tf.keras.layers.Dense(10, activation='softmax')
])
训练模型
model.compile(optimizer='adam',
loss='categorical_crossentropy',
metrics=['accuracy'])
model.fit(x_train, y_train, epochs=10)
测试模型
test_loss, test_acc = model.evaluate(x_test, y_test)
print('Test accuracy:', test_acc)
import matplotlib.pyplot as plt
import numpy as np
predictions = model.predict(x_test)
# 随机选择一些测试图像
indices = np.random.choice(range(len(x_test)), 10)
predictions = model.predict(x_test)
fig, axs = plt.subplots(2,5, figsize=(20,8))
# 可视化测试图像及其预测标签
for i, ax in zip(indices, axs.flatten()):
ax.imshow(x_test[i], cmap='gray')
ax.set_title(f"Predicted label: {np.argmax(predictions[i])}")
plt.show()
Pytorch版本
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
# 加载MNIST数据集
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)
# 定义神经网络模型
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(28*28, 128)
self.fc2 = nn.Linear(128, 64)
self.fc3 = nn.Linear(64, 10)
def forward(self, x):
x = x.view(-1, 28*28) # 将图像展平为向量
x = torch.relu(self.fc1(x))
x = torch.relu(self.fc2(x))
x = torch.softmax(self.fc3(x), dim=1) # 使用softmax输出概率分布
return x
net = Net()
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss() # 注意:CrossEntropyLoss内部进行了log_softmax操作,因此输出层不需要再softmax
optimizer = optim.Adam(net.parameters(), lr=0.001)
# 训练模型
for epoch in range(10): # 迭代10个epoch
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs, labels = data
optimizer.zero_grad() # 清空梯度
outputs = net(inputs) # 前向传播
loss = criterion(outputs, labels) # 计算损失
loss.backward() # 反向传播
optimizer.step() # 更新参数
running_loss += loss.item()
print(f'Epoch [{epoch+1}/10], Loss: {running_loss/len(trainloader):.4f}')
# 在测试集上评估模型
correct = 0
total = 0
with torch.no_grad(): # 评估模式,不需要计算梯度
for data in testloader:
images, labels = data
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Accuracy of the network on the 10000 test images: {100 * correct / total:.2f}%')
# 可视化测试图像及其预测标签
predictions = []
test_images, test_labels = next(iter(testloader)) # 一次性加载整个测试集可能会占用大量内存,这里只取一个batch
with torch.no_grad():
test_outputs = net(test_images)
_, predicted_labels = torch.max(test_outputs, 1)
predictions.append(predicted_labels.numpy())
predictions = np.concatenate(predictions) # 虽然这里只有一个batch,但为了与TensorFlow代码风格一致,仍然使用concatenate
indices = np.random.choice(range(len(test_images)), 10)
fig, axs = plt.subplots(2, 5, figsize=(20, 8))
for i, ax in zip(indices, axs.flatten()):
ax.imshow(test_images[i].squeeze().numpy(), cmap='gray') # 转换回numpy数组并去除多余的维度
ax.set_title(f"Predicted label: {predictions[i]}")
plt.show()
原文地址:https://blog.csdn.net/Zsusan7/article/details/142734421
免责声明:本站文章内容转载自网络资源,如本站内容侵犯了原著者的合法权益,可联系本站删除。更多内容请关注自学内容网(zxcms.com)!