知识点回顾：
- PyTorch和cuda的安装
- 查看显卡信息的命令行命令（cmd中使用）
- cuda的检查
- 简单神经网络的流程
- 数据预处理（归一化、转换成张量）
- 模型的定义
- 继承nn.Module类
- 定义每一个层
- 定义前向传播流程
- 定义损失函数和优化器
- 定义训练流程
- 可视化loss过程
# test_pytorch.py import torch # 检查导入是否成功 print(f"PyTorch版本: {torch.__version__}") print(f"CUDA可用: {torch.cuda.is_available()}") # 设置随机种子 torch.manual_seed(42) # 创建一个张量测试 x = torch.randn(5, 5) print(f"随机张量:\n{x}") import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset import numpy as np import matplotlib.pyplot as plt # 设置随机种子确保结果可复现 torch.manual_seed(42) np.random.seed(42) # ------------------------- # 1. 环境检查 # ------------------------- def check_environment(): """检查CUDA可用性并打印环境信息""" if torch.cuda.is_available(): print(f"CUDA可用 - 设备: {torch.cuda.get_device_name(0)}") print(f"CUDA版本: {torch.version.cuda}") else: print("CUDA不可用，将使用CPU进行计算") print(f"PyTorch版本: {torch.__version__}") # ------------------------- # 2. 数据预处理 # ------------------------- class DataPreprocessor: """数据预处理类，支持回归和分类任务""" def __init__(self, task_type="regression"): """ task_type: 'regression' 或 'classification' """ self.task_type = task_type self.x_mean, self.x_std = None, None self.y_mean, self.y_std = None, None def preprocess(self, x, y): """预处理数据并转换为张量""" # 特征归一化 self.x_mean, self.x_std = x.mean(), x.std() x_normalized = (x - self.x_mean) / self.x_std # 根据任务类型处理标签 if self.task_type == "regression": self.y_mean, self.y_std = y.mean(), y.std() y_normalized = (y - self.y_mean) / self.y_std y_tensor = torch.tensor(y_normalized, dtype=torch.float32).view(-1, 1) else: # 分类任务 y_tensor = torch.tensor(y, dtype=torch.long) x_tensor = torch.tensor(x_normalized, dtype=torch.float32).view(-1, 1) return x_tensor, y_tensor def denormalize_prediction(self, pred): """将预测结果反归一化（仅用于回归任务）""" if self.task_type == "regression": return pred * self.y_std + self.y_mean return pred # ------------------------- # 3. 模型定义 # ------------------------- class SimpleNet(nn.Module): """简单的两层神经网络""" def __init__(self, input_size=1, hidden_size=10, output_size=1): super(SimpleNet, self).__init__() self.fc1 = nn.Linear(input_size, hidden_size) self.relu = nn.ReLU() self.fc2 = nn.Linear(hidden_size, output_size) def forward(self, x): x = self.fc1(x) x = self.relu(x) x = self.fc2(x) return x # ------------------------- # 4. 训练模块 # ------------------------- class Trainer: """模型训练器""" def __init__(self, model, task_type="regression", lr=0.01): self.model = model self.task_type = task_type # 根据任务类型选择损失函数 self.criterion = nn.MSELoss() if task_type == "regression" else nn.CrossEntropyLoss() self.optimizer = optim.Adam(model.parameters(), lr=lr) self.train_losses = [] def train(self, dataloader, num_epochs=100, print_every=10): """训练模型并记录损失""" for epoch in range(num_epochs): epoch_loss = 0.0 for inputs, targets in dataloader: # 前向传播 outputs = self.model(inputs) loss = self.criterion(outputs, targets) # 反向传播和优化 self.optimizer.zero_grad() loss.backward() self.optimizer.step() epoch_loss += loss.item() # 记录每个epoch的平均损失 avg_loss = epoch_loss / len(dataloader) self.train_losses.append(avg_loss) if (epoch + 1) % print_every == 0: print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {avg_loss:.4f}') return self.train_losses def plot_loss(self): """可视化训练损失""" plt.figure(figsize=(10, 6)) plt.plot(range(1, len(self.train_losses) + 1), self.train_losses, label='Training Loss') plt.xlabel('Epoch') plt.ylabel('Loss') plt.title('Training Loss over Epochs') plt.legend() plt.grid(True) plt.show() # ------------------------- # 5. 主函数 # ------------------------- def main(): # 检查环境 check_environment() # 生成示例数据 task_type = "regression" # 或 "classification" if task_type == "regression": # 回归任务数据 x = np.linspace(0, 10, 100) y = 2 * x + 1 + np.random.randn(100) * 0.5 # y = 2x + 1 + 噪声 else: # 分类任务数据 x = np.linspace(0, 10, 100) y = (x > 5).astype(int) # 二分类问题 # 数据预处理 preprocessor = DataPreprocessor(task_type) x_tensor, y_tensor = preprocessor.preprocess(x, y) # 创建数据加载器 dataset = TensorDataset(x_tensor, y_tensor) dataloader = DataLoader(dataset, batch_size=16, shuffle=True) # 初始化模型 if task_type == "regression": model = SimpleNet(input_size=1, output_size=1) else: model = SimpleNet(input_size=1, output_size=2) # 二分类问题有2个类别 # 训练模型 trainer = Trainer(model, task_type=task_type, lr=0.01) losses = trainer.train(dataloader, num_epochs=100) # 可视化损失 trainer.plot_loss() # 评估模型 model.eval() with torch.no_grad(): test_x = np.array([5.0]) # 测试输入 test_x_normalized = (test_x - preprocessor.x_mean) / preprocessor.x_std test_tensor = torch.tensor(test_x_normalized, dtype=torch.float32).view(-1, 1) prediction = model(test_tensor) if task_type == "regression": # 反归一化回归预测结果 prediction = preprocessor.denormalize_prediction(prediction.item()) print(f"回归预测结果: {prediction:.4f}") print(f"真实值: {2*test_x[0] + 1:.4f}") # 基于y=2x+1 else: # 获取分类预测结果 _, predicted = torch.max(prediction, 1) print(f"分类预测类别: {predicted.item()}") print(f"真实类别: {1 if test_x[0] > 5 else 0}") if __name__ == "__main__": main()