import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

运行结果

 Using device: cuda 

inputs = torch.rand((100,3)) # 100 条数据，3个维度， 都是0-1之间的随机数
# 设置权重
weights = torch.tensor([[1.1],[2.2],[3.3]]) # 3行1列
bias = torch.tensor([4.4]) # 偏置
targets = inputs @ weights + bias + 0.1*torch.rand((100,1))  # 100行1列

# 初始化，启用在梯度
w = torch.randn((3,1), requires_grad=True,device=device)
b = torch.randn(1, requires_grad=True,device=device)
lr = 0.01 # 学习率
epoch = 10000 # 迭代次数

# 确保 inputs 和 targets 在同一设备上
inputs = inputs.to(device)
targets = targets.to(device)

for i in range(epoch):
    outputs = inputs @ w + b
    loss = torch.mean((outputs - targets)**2) # 均方误差
    if i % 1000 == 0:
        print(i, loss.item())
    loss.backward() # 反向传播，计算梯度
    with torch.no_grad(): # 关闭梯度计算
        w -= lr * w.grad # 更新权重
        b -= lr * b.grad # 更新偏置
        w.grad.zero_() # 清零梯度
        b.grad.zero_() # 清零梯度

print('预测的权重：', w)
print('预测的偏置：', b)

运行结果

 预测的权重： tensor([[1.0846], [2.1862], [3.2870]], device='cuda:0', requires_grad=True) 预测的偏置： tensor([4.4735], device='cuda:0', requires_grad=True)