# 1. 查看當前的芯片 CPU 資源分配情況
# cpu-num-cfg-i 0-c 0 表示查詢 ID 為 0 的芯片的 CPU 配置
npu-smi info-t cpu-num-cfg-i 0-c 0
# 2. 調整 CPU 配比
# 格式解釋：AI CPU : Control CPU : Data CPU
# 設置為 3 : 1 : 0
sudo npu-smi set-t cpu-num-cfg-i 0-c 0-v 3:1:0

"""
CANN × 智慧農業：植物病蟲害識別模型訓練腳本
"""
import os
import random
import datetime
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset
from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights
from tqdm import tqdm
import matplotlib.pyplot as plt
# ========== 參數配置 ==========
data_dir = "./datasets/PlantVillage/PlantVillage"
batch_size = 32
num_epochs = 25
lr = 0.001
train_ratio = 0.8
# ========== 設備優先級：NPU > CPU ==========
try:
import torch_npu  # noqa: F401
npu_available = torch.npu.is_available()
except Exception:
npu_available = False
if npu_available:
device = torch.device("npu:0")
torch.npu.set_device(device)
else:
device = torch.device("cpu")
print(f"使用設備: {device}")
# ========== 數據增強 ==========
train_transform = transforms.Compose([
transforms.Resize((256, 256)),
transforms.RandomCrop(224),
transforms.RandomHorizontalFlip(0.5),
transforms.RandomRotation(15),
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
val_transform = transforms.Compose([
transforms.Resize((256, 256)),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# ========== 加載數據集並劃分 ==========
print("加載數據集...")
train_dataset = datasets.ImageFolder(data_dir, transform=train_transform)
val_dataset = datasets.ImageFolder(data_dir, transform=val_transform)
num_classes = len(train_dataset.classes)
print(f"類別數: {num_classes}")
print(f"總樣本數: {len(train_dataset)}")
indices = list(range(len(train_dataset)))
random.seed(42)
random.shuffle(indices)
split = int(train_ratio * len(indices))
train_indices = indices[:split]
val_indices = indices[split:]
train_subset = Subset(train_dataset, train_indices)
val_subset = Subset(val_dataset, val_indices)
os.makedirs("output", exist_ok=True)
train_loader = DataLoader(train_subset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
val_loader = DataLoader(val_subset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)
print(f"訓練集樣本數: {len(train_subset)}")
print(f"驗證集樣本數: {len(val_subset)}")
# ========== 加載模型 ==========
print("\n加載 EfficientNet-B0 模型...")
try:
weights = EfficientNet_B0_Weights.DEFAULT
model = efficientnet_b0(weights=weights)
print("✅ 成功加載預訓練權重")
except Exception as e:
print(f"⚠️ 加載預訓練權重失敗: {e}")
model = efficientnet_b0(weights=None)
# 凍結特徵層，只訓練分類頭
for p in model.features.parameters():
p.requires_grad = False
# 替換分類頭
in_features = model.classifier[1].in_features
model.classifier[1] = nn.Linear(in_features, num_classes)
model = model.to(device)
# 統計參數
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"總參數數: {total_params:,}")
print(f"可訓練參數數: {trainable_params:,}")
# ========== 優化配置 ==========
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam((p for p in model.parameters() if p.requires_grad), lr=lr)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=5, verbose=True)
# ========== 訓練/驗證 ==========
def train_epoch(model, loader, criterion, optimizer, device):
model.train()
running_loss, correct, total = 0.0, 0, 0
for imgs, labels in tqdm(loader, desc="Training", leave=False):
imgs, labels = imgs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(imgs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item() * imgs.size(0)
_, pred = outputs.max(1)
total += labels.size(0)
correct += pred.eq(labels).sum().item()
return running_loss / total, 100.0 * correct / total
def validate(model, loader, criterion, device):
model.eval（)
running_loss, correct, total = 0.0, 0, 0
with torch.no_grad():
for imgs, labels in tqdm(loader, desc="Validating", leave=False):
imgs, labels = imgs.to(device), labels.to(device)
outputs = model(imgs)
loss = criterion(outputs, labels)
running_loss += loss.item() * imgs.size(0)
_, pred = outputs.max(1)
total += labels.size(0)
correct += pred.eq(labels).sum().item()
return running_loss / total, 100.0 * correct / total
# ========== 訓練主循環 ==========
print("\n開始訓練...")
print(f"開始時間: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
train_losses, val_losses, train_accs, val_accs = [], [], [], []
best_val_acc, best_model_path = 0.0, None
start_time = datetime.datetime.now()
for epoch in range(num_epochs):
print(f"\nEpoch {epoch + 1}/{num_epochs}")
print("-" * 50)
tr_loss, tr_acc = train_epoch(model, train_loader, criterion, optimizer, device)
va_loss, va_acc = validate(model, val_loader, criterion, device)
train_losses.append(tr_loss); train_accs.append(tr_acc)
val_losses.append(va_loss); val_accs.append(va_acc)
scheduler.step(va_acc)
print(f"Train Loss: {tr_loss:.4f}, Train Acc: {tr_acc:.2f}%")
print(f"Val   Loss: {va_loss:.4f}, Val   Acc: {va_acc:.2f}%")
if va_acc > best_val_acc:
best_val_acc = va_acc
best_model_path = os.path.join("output", "best_plant_disease_model.pth")
torch.save(model.state_dict(), best_model_path)
torch.save(model.state_dict(), "best_plant_disease_model.pth")  # 兼容其它腳本
print(f"✅ 保存最佳模型，驗證準確率: {va_acc:.2f}%")
if (epoch + 1) % 5 == 0:
ckpt = os.path.join("output", f"checkpoint_epoch_{epoch + 1}.pth")
torch.save(model.state_dict(), ckpt)
print(f"✅ 中間檢查點已保存: {ckpt}")
if best_model_path is None:
best_model_path = os.path.join("output", "best_plant_disease_model.pth")
torch.save(model.state_dict(), best_model_path)
torch.save(model.state_dict(), "best_plant_disease_model.pth")
end_time = datetime.datetime.now()
print(f"\n訓練完成時間: {end_time.strftime('%Y-%m-%d %H:%M:%S')}")
print(f"訓練總時長: {end_time - start_time}")
# ========== 繪製訓練曲線 ==========
plt.figure(figsize=(14, 5))
plt.subplot(1, 2, 1)
plt.plot(train_losses, label='Train Loss', marker='o')
plt.plot(val_losses, label='Val Loss', marker='s')
plt.xlabel('Epoch'); plt.ylabel('Loss'); plt.legend(); plt.title('Loss'); plt.grid(True, alpha=0.3)
plt.subplot(1, 2, 2)
plt.plot(train_accs, label='Train Acc', marker='o')
plt.plot(val_accs, label='Val Acc', marker='s')
plt.xlabel('Epoch'); plt.ylabel('Accuracy (%)'); plt.legend(); plt.title('Accuracy'); plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig(os.path.join('output', 'training_curves.png'), dpi=300, bbox_inches='tight')
plt.savefig('training_curves.png', dpi=300, bbox_inches='tight')
print("✅ 訓練曲線已保存: output/training_curves.png, training_curves.png")
print("\n" + "=" * 60)
print("訓練完成！")
print("最佳驗證準確率: %.2f%%" % best_val_acc)
print("模型路徑: %s" % best_model_path)
print("=" * 60)

# 定義虛擬輸入（Dummy Input），用於觸發一次前向傳播，從而追蹤網絡結構
# 必須與實際推理時的分辨率一致 (1, 3, 224, 224)
dummy_input = torch.randn(1, 3, 224, 224).to(device)
torch.onnx.export(model,
dummy_input,
"plant_disease.onnx",
input_names=["input"],
output_names=["output"],
# opset_version=11 是目前兼容性最好的版本之一
opset_version=11)

atc --model=plant_disease.onnx \
--framework=5 \                  # 5 代表 ONNX 框架
--output=plant_disease_cann \    # 輸出文件名
--input_format=NCHW \
--input_shape="input:1,3,224,224" \
--soc_version=Ascend310B1 \      # 指定目標芯片型號，針對特定架構優化指令
--log=info                       # 開啓日誌，可以看到算子融合的詳細過程