同时搞定TensorFlow、PyTorch (四)：模型训练

同时搞定TensorFlow、PyTorch (一)：梯度下降。同时搞定TensorFlow、PyTorch (二)：模型定义。同时搞定TensorFlow、PyTorch (三) ：资料前置处理。同时搞定TensorFlow、PyTorch (四)：模型训练。

前言

上一篇谈到TensorFlow/PyTorch资料前置处理的方式，接下来我们就各使用一支完整的範例训练模型。

TensorFlow

载入 MNIST 资料集。

import tensorflow as tfmnist = tf.keras.datasets.mnist# 载入 MNIST 手写阿拉伯数字资料(x_train, y_train),(x_test, y_test) = mnist.load_data()# 特徵缩放，使用常态化(Normalization)，公式 = (x - min) / (max - min)x_train_norm, x_test_norm = x_train / 255.0, x_test / 255.0# 转为 Dataset，含 X/Y 资料train_ds = tf.data.Dataset.from_tensor_slices((x_train_norm, y_train))

建立模型。

model = tf.keras.models.Sequential([  tf.keras.layers.Flatten(input_shape=(28, 28)),  tf.keras.layers.Dense(128, activation='relu'),  tf.keras.layers.Dropout(0.2),  tf.keras.layers.Dense(10, activation='softmax')])

设定优化器(optimizer)、损失函数(loss)、效能衡量指标(metrics)。

model.compile(optimizer='adam',              loss='sparse_categorical_crossentropy',              metrics=['accuracy'])

损失函数使用sparse_XXX，y就不需One-hot encoding，否则须呼叫以下程式。

# 将 training 的 label 进行 one-hot encoding，例如数字 7 经过 One-hot encoding 转换后是 0000000100，即第8个值为 1y_train = tf.keras.utils.to_categorical(y_train) y_test = tf.keras.utils.to_categorical(y_test) model.compile(optimizer='adam',              loss='categorical_crossentropy',              metrics=['accuracy'])

模型训练：只要一行程式码。

model.fit(x_train_norm, y_train, epochs=5, validation_split=0.2)

评分(Score Model)。

score=model.evaluate(x_test_norm, y_test, verbose=0)for i, x in enumerate(score):    print(f'{model.metrics_names[i]}: {score[i]:.4f}')

PyTorch

载入 MNIST 资料集。

import osimport torchfrom torch import nnfrom torch.nn import functional as Ffrom torch.utils.data import DataLoader, random_splitfrom torchmetrics import Accuracyfrom torchvision import transformsfrom torchvision.datasets import MNISTPATH_DATASETS = "" # 预设路径# 下载 MNIST 手写阿拉伯数字 训练资料train_ds = MNIST(PATH_DATASETS, train=True, download=True,                  transform=transforms.ToTensor())# 下载测试资料test_ds = MNIST(PATH_DATASETS, train=False, download=True,                  transform=transforms.ToTensor())

判断本机是否有GPU：若含NVidia显卡，且安装CUDA toolkits，TensorFlow会自动使用GPU进行运算及模型训练，PyTorch则由开发者自行决定，因此，先判断本机是否有GPU。

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

建立模型：若有GPU，to(device)指令将模型複製至GPU记忆体。

model = torch.nn.Sequential(    torch.nn.Flatten(),    torch.nn.Linear(28 * 28, 256),     torch.nn.Dropout(0.2),    torch.nn.Linear(256, 10),     # 使用nn.CrossEntropyLoss()时，不需要将输出经过softmax层，否则计算的损失会有误    # torch.nn.Softmax(dim=1)).to(device)

设定优化器(optimizer)、损失函数(loss)、效能衡量指标(metrics)。

epochs = 5lr=0.1BATCH_SIZE = 1024  # 批量# 建立 DataLoadertrain_loader = DataLoader(train_ds, batch_size=600)# 设定优化器(optimizer)# optimizer = torch.optim.Adam(model.parameters(), lr=lr)optimizer = torch.optim.Adadelta(model.parameters(), lr=lr)criterion = nn.CrossEntropyLoss()

PyTorch不管使用何种损失函数，y是否进行One-hot encoding均可。

模型训练：要写一大堆程式码，包括梯度下降、计算损失/準确率等，似乎很麻烦，但是训练过程中，可任意插入除错讯息，例如第6行，显示第一笔资料，不像TensorFlow必须借助Callback机制。另外，要记得把训练资料也複製到GPU记忆体，与模型一致。

model.train()loss_list = []    for epoch in range(1, epochs + 1):    for batch_idx, (data, target) in enumerate(train_loader):        data, target = data.to(device), target.to(device)#         if batch_idx == 0 and epoch == 1: print(data[0])        optimizer.zero_grad()        output = model(data)        loss = criterion(output, target)        loss.backward()        optimizer.step()        if batch_idx % 10 == 0:            loss_list.append(loss.item())            batch = batch_idx * len(data)            data_count = len(train_loader.dataset)            percentage = (100. * batch_idx / len(train_loader))            print(f'Epoch {epoch}: [{batch:5d} / {data_count}] ({percentage:.0f} %)' +                  f'  Loss: {loss.item():.6f}')

评分(Score Model)：也要写一大堆程式码。

# 建立 DataLoadertest_loader = DataLoader(test_ds, shuffle=False, batch_size=BATCH_SIZE)model.eval()test_loss = 0correct = 0with torch.no_grad():    for data, target in test_loader:        data, target = data.to(device), target.to(device)        output = model(data)                # sum up batch loss        test_loss += criterion(output, target).item()                # 预测        pred = output.argmax(dim=1, keepdim=True)                  # 正确笔数        correct += pred.eq(target.view_as(pred)).sum().item()# 平均损失test_loss /= len(test_loader.dataset)# 显示测试结果batch = batch_idx * len(data)data_count = len(test_loader.dataset)percentage = 100. * correct / data_countprint(f'平均损失: {test_loss:.4f}, 準确率: {correct}/{data_count}' +       f' ({percentage:.0f}%)\n')

完整程式可参阅开发者传授 PyTorch 秘笈 的src/TF_vs_Pytorch.ipynb。

详细流程强烈建议参考：

TensorFlow：深度学习 -- 最佳入门迈向 AI 专题实战的src/04_02_手写阿拉伯数字辨识_完整版。PyTorch可参阅开发者传授 PyTorch 秘笈 的src/04_05_手写阿拉伯数字辨识_完整版.ipynb。

结论

TensorFlow训练及评分都只要一行指令，PyTorch 可以吗? 可以试用 PyTorch Lightning 套件，首页上附了一支简短的程式，训练也只要一行指令，读者如有兴趣可详阅开发者传授 PyTorch 秘笈。PyTorch 训练及评分步骤比较繁琐，但是相对程序较清楚，学界会比较喜爱PyTorch，因为，可以讲授得很清楚，要调整也比较容易，因此，PyTorch在研发及教育领域佔有率比TensorFlow高，反之，企业界讲求效率，使用TensorFlow较多。

系列文介绍到此告一段落，希望对深度学习入门者有一丝丝的帮助。

以下为工商广告:)。
PyTorch：
开发者传授 PyTorch 秘笈

预计 2022/6/20 出版。

TensorFlow：
深度学习 -- 最佳入门迈向 AI 专题实战。