?一、第三方庫介紹

二、訓練數據集介紹

三、原理簡介

????????該代碼使用PyTorch訓練一個基于ResNet-18的貓狗分類模型。通過加載并處理數據、訓練模型、調整最后輸出層、使用Adam優化器進行反向傳播，并在每個訓練周期輸出損失與準確率。訓練完畢后，保存模型用于后續預測。

四、代碼思路簡介

• 加載數據?→ 使用?datasets.ImageFolder?加載貓狗數據集，并應用圖像轉換。
• 構建模型?→ 使用預訓練的 ResNet-18 模型，修改輸出層以適應2類分類。
• 定義損失和優化器?→ 使用交叉熵損失函數和 Adam 優化器。
• 訓練模型?→ 遍歷數據集，前向傳播、計算損失、反向傳播、更新模型參數。
• 保存模型?→ 訓練完成后，保存模型權重。
• 預測圖片?→ 加載已訓練模型，輸入圖片進行預測并輸出分類結果。

五、代碼

場景：使用pytorch識別貓狗
貓的圖片路徑：F:\pycharm\AIDEMO\data\cat
狗的圖片路徑：F:\pycharm\AIDEMO\data\dog
需要判斷的圖片：F:\pycharm\AIDEMO\01.jpeg

import torch
import torchvision
from torchvision import datasets, transforms
from torch import nn, optim
from torch.utils.data import DataLoader
from PIL import Image
from torchvision.models import ResNet18_Weights# 定義transform類（視覺轉換類，將圖片格式轉化為張量格式）
transform = transforms.Compose([transforms.Resize((128, 128)),  # 將圖片縮放到統一大小transforms.ToTensor(),  # 轉換為Tensor格式transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 標準化處理
])def train_model(data_dir, num_epochs=10, batch_size=32, save_path='cat_dog_model.pth'):"""訓練模型并保存。:param data_dir: 數據路徑，包含cat和dog文件夾:param num_epochs: 訓練周期，默認為10:param batch_size: 批次大小，默認為32:param save_path: 模型保存路徑，默認為'cat_dog_model.pth'"""# 1. 加載訓練數據train_data = datasets.ImageFolder(root=data_dir,  # 數據路徑transform=transform)print(train_data.class_to_idx)  # 輸出文件夾編號，例如這里輸出{'cat': 0, 'dog': 1}，表達0代表貓貓，1代表狗狗train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)# 2. 使用預訓練的ResNet18模型model = torchvision.models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)model.fc = nn.Linear(model.fc.in_features, 2)  # 修改輸出層以適應2類分類（貓、狗）# 3. 定義損失函數和優化器criterion = nn.CrossEntropyLoss()optimizer = optim.Adam(model.parameters(), lr=0.001)# 4. 開始訓練模型for epoch in range(num_epochs):model.train()  # 設置模型為訓練模式running_loss = 0.0  # 初始化損失值correct = 0  # 模型預測準確數total = 0  # 模型預測總數for images, labels in train_loader:optimizer.zero_grad()  # 清除之前的梯度outputs = model(images)  # 前向傳播，得出預測結果loss = criterion(outputs, labels)  # 計算損失loss.backward()  # 反向傳播optimizer.step()  # 更新參數running_loss += loss.item()# 計算準確率_, predicted = torch.max(outputs, 1)  # 獲取預測結果total += labels.size(0)  # 累計總樣本數correct += (predicted == labels).sum().item()  # 累計預測正確的樣本數# 輸出訓練周期的損失和準確率accuracy = 100 * correct / total  # 計算準確率print(f'周期 [{epoch + 1}/{num_epochs}], 損失: {running_loss / len(train_loader):.4f}, 準確率: {accuracy:.2f}%')if accuracy == 100:  # 準確率達到100%就停止訓練，避免過度擬合break# 保存訓練模型torch.save(model.state_dict(), save_path)def predict_image(model_path, img_path):"""加載訓練好的模型并進行圖片預測。:param model_path: 訓練好的模型路徑:param img_path: 需要預測的圖片路徑:return: 預測結果（貓或狗）"""# 加載模型model = torchvision.models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)model.fc = nn.Linear(model.fc.in_features, 2)model.load_state_dict(torch.load(model_path))model.eval()  # 設置模型為評估模式# 預測指定圖片img = Image.open(img_path)img = transform(img).unsqueeze(0)  # 將圖片處理成張量輸出并增加batch維度# 模型預測with torch.no_grad():  # 不需要梯度計算，只是進行模型預測outputs = model(img)_, predicted = torch.max(outputs, 1)# 輸出預測結果return "這是貓的圖片" if predicted.item() == 0 else "這是狗的圖片"if __name__ == "__main__":# 01 訓練出模型（若已訓練出準確度較高模型，可注釋下面兩句話，直接用訓練完畢的模型預測）data_dir = 'F:/pycharm/AIDEMO/data'  # 數據路徑train_model(data_dir, num_epochs=10, batch_size=32)# 02 預測指定圖片img_path = 'F:/pycharm/AIDEMO/data/01.jpeg'  # 圖片路徑result = predict_image('cat_dog_model.pth', img_path)print(result)

六、輸出結果展示