pytorch y=2x学习,写这个模型,并且运行一个测试判断模型训练效果
时间: 2023-04-01 08:00:47 浏览: 73
可以使用以下代码实现:
import torch
# 定义模型
class Model(torch.nn.Module):
def __init__(self):
super(Model, self).__init__()
self.linear = torch.nn.Linear(1, 1)
def forward(self, x):
y_pred = self.linear(x)
return y_pred
# 定义数据
x_train = torch.tensor([[1.], [2.], [3.], [4.]])
y_train = torch.tensor([[2.], [4.], [6.], [8.]])
# 定义损失函数和优化器
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=.01)
# 训练模型
model = Model()
for epoch in range(100):
y_pred = model(x_train)
loss = criterion(y_pred, y_train)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 测试模型
x_test = torch.tensor([[5.]])
y_test = model(x_test)
print(y_test) # 输出预测结果
注意:以上代码仅供参考,实际应用中可能需要根据具体情况进行修改。