pytorch深度學習分類程式碼簡單示例

train.py程式碼如下

import torch
import torch.nn as nn
import torch.optim as optim

model_save_path = "my_model.pth"

# 定義簡單的線性神經網路模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.output = nn.Linear(2, 4)   # 輸入2個特徵，輸出4個類別

    def forward(self, x):
        x = self.output(x)
        return x

def main():
    # 資料點
    x = torch.tensor([[0, 0], [0, 10], [10, 0], [10, 10]], dtype=torch.float32)
    y = torch.tensor([0, 1, 2, 3], dtype=torch.long)

    # 初始化模型
    model = MyModel()

    # 定義損失函式和最佳化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.01)

    # 訓練模型
    num_iterations = 10000  # 迭代次數
    for i in range(num_iterations):
        model.train()

        # 前向傳播：計算預測輸出
        y_pred = model(x)

        # 計算損失
        loss = criterion(y_pred, y)

        # 輸出每1000次迭代的損失值
        if i % 1000 == 0:
            print(f"迭代 {i}，損失：{loss.item():.4f}")

        # 反向傳播與梯度更新
        optimizer.zero_grad()  # 清除梯度
        loss.backward()  # 計算梯度
        optimizer.step()  # 更新引數

    # 列印最佳化後的權重和偏置
    print("\n最佳化後的權重和偏置:")
    for name, param in model.named_parameters():
        if param.requires_grad:
            print(f"{name} = {param.data.numpy()}")

    # 儲存模型
    torch.save(model.state_dict(), model_save_path)
    print(f"模型已儲存到 {model_save_path}")

if __name__ == "__main__":
    main()

執行結果

test.py程式碼如下

import numpy as np
import torch
from torch import nn

from train import MyModel, model_save_path

# 載入模型
loaded_model = MyModel()
loaded_model.load_state_dict(torch.load(model_save_path))
loaded_model.eval()  # 切換到評估模式

# 定義預測資料
input_data = [0, 9]

# 使用載入的模型進行預測
x_new = torch.tensor([input_data], dtype=torch.float32)  # 新資料點
y_new_pred = loaded_model(x_new)  # 計算預測值

# 使用softmax計算每個類別的機率
softmax = nn.Softmax(dim=1)
y_new_pred_probs = softmax(y_new_pred)

# 找到預測的類別
predicted_class = torch.argmax(y_new_pred_probs, dim=1)

# 將機率分佈四捨五入到三位小數
y_new_pred_probs_rounded = np.round(y_new_pred_probs.detach().numpy(), 3)

print(f"\n對x = {input_data}的預測類別：{predicted_class.item()}")
print(f"預測類別的機率分佈：{y_new_pred_probs_rounded}")

# 列印權重和偏置
weights = loaded_model.output.weight  # 獲取輸出層權重
bias = loaded_model.output.bias  # 獲取輸出層偏置

print(f"\n模型權重：\n{weights}")
print(f"\n模型偏置：\n{bias}")

# 計算input_data * 模型權重 + 模型偏置
with torch.no_grad():
    linear_output = x_new @ weights.t() + bias

print(f"\ninput_data * weights + bias ={linear_output.numpy()}")

# 手動計算Softmax機率分佈
linear_output_np = linear_output.numpy()
exp_output = np.exp(linear_output_np)
softmax_output_manual = exp_output / np.sum(exp_output)

print(f"\n手動計算的Softmax機率分佈：{softmax_output_manual}")
print(f"手動計算的預測類別：{np.argmax(softmax_output_manual)}")

執行結果