tian/CostPrediction
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: 缺PyTorch时训练不会因ImportError中断整个流程

Browse Source
This commit is contained in:
tian 2026-06-10 09:41:19 +08:00
parent b6b072d844
commit 531cff509d
1 changed files with 9 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/model_trainer.py
View File

@ -500,12 +500,16 @@ class ModelTrainer:
model.fit(X_train, y_train) model.fit(X_train, y_train)
elif model_type == 'pytorch': elif model_type == 'pytorch':
# 训练PyTorch模型 # 训练PyTorch模型（如未安装则跳过）
try:
train_dataset = EquipmentDataset(X_train, y_train) train_dataset = EquipmentDataset(X_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True) train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
training_success = self.train_model(train_loader) training_success = self.train_model(train_loader)
if not training_success: if not training_success:
continue continue
except ImportError as e:
logger.warning(f"PyTorch not available, skipping: {e}")
continue
# 评估模型性能 # 评估模型性能
if model_type == 'pytorch': if model_type == 'pytorch':