41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
import torch
|
||
|
||
# 检查MPS可用性（需要PyTorch 1.12+和macOS 12.3+）
|
||
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
|
||
|
||
# 生成训练数据（移动到MPS设备）
|
||
X = torch.randn(1000, 2).to(device) # 1000个样本，2个特征
|
||
y = X @ torch.tensor([2.0, -3.4], device=device) + 4 # 真实关系式
|
||
y += 0.01 * torch.randn(y.shape, device=device) # 添加噪声
|
||
|
||
# 定义模型（必须继承nn.Module）
|
||
class LinearRegression(torch.nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.linear = torch.nn.Linear(2, 1) # 输入2维，输出1维
|
||
|
||
def forward(self, x):
|
||
return self.linear(x)
|
||
|
||
model = LinearRegression().to(device) # 将模型移至MPS设备
|
||
criterion = torch.nn.MSELoss()
|
||
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
|
||
|
||
# 训练循环
|
||
for epoch in range(500):
|
||
# 前向传播
|
||
outputs = model(X)
|
||
loss = criterion(outputs, y.unsqueeze(1))
|
||
|
||
# 反向传播
|
||
optimizer.zero_grad()
|
||
loss.backward()
|
||
optimizer.step()
|
||
|
||
if epoch % 50 == 0:
|
||
print(f'Epoch {epoch}, loss: {loss.item():.4f}')
|
||
|
||
# 输出最终参数
|
||
print("Learned weights:", model.linear.weight.data)
|
||
print("Learned bias:", model.linear.bias.data)
|