torch训练简单例子
import torch
import torch.nn as nn
import torch.optim as optim
# 定义一个简单的模型
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.fc = nn.Linear(10, 5)
def forward(self, x):
return self.fc(x)
# 创建模型实例
model = MyModel()
# 检查GPU是否可用
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 将模型移动到GPU
model.to(device)
# 创建优化器
optimizer = optim.SGD(model.parameters(), lr=0.1)
# 定义随机数生成的样本数量
batch_size = 32
# 定义训练轮数
num_epochs = 10
# 执行训练循环
for epoch in range(num_epochs):
model.train() # 设置模型为训练模式
for _ in range(batch_size):
# 生成随机输入数据和目标数据
inputs = torch.randn(batch_size, 10).to(device)
targets = torch.randn(batch_size, 5).to(device)
# 前向传播
outputs = model(inputs)
loss = nn.MSELoss()(outputs, targets)
# 反向传播和梯度更新
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 在评估阶段禁用梯度计算
with torch.no_grad():
model.eval() # 设置模型为评估模式
for _ in range(batch_size):
# 生成随机输入数据
inputs = torch.randn(batch_size, 10).to(device)
# 前向传播
outputs = model(inputs)
# 对输出进行处理...