pytorch神经网络训练
2023-03-28
技术
学自李宏毅老师
神经网络训练设置
# 从数据集中读取数据
dataset = MyDataset(file)
# 将数据放到数据加载器,分批
tr_set = DataLoader(dataset, 16, shuffle=True)
# 构造模型并指定设备(cpu/cuda(gpu))
model = MyModel().to(device)
# 设置损失函数
criterion = nn.MSELoss()
# 设置优化器
optimizer = torch.optim.SGD(model.parameters(), 0.1)
神经网络训练回路
感觉这里用回路比循环好听一点
# 迭代 n_epochs
for epoch in range(n_epochs):
# 将模型设置为训练模式
model.train()
# 通过数据加载器迭代
for x, y in tr_set:
# 梯度置为 0
optimizer.zero_grad()
# cpu/cuda
x, y = x.to(device), y.to(device)
# 正向传播(计算输出)
pred = model(x)
# 计算损失
loss = criterion(pred, y)
# 反向传播(计算梯度)
loss.backward()
# 使用优化器更新模型
optimizer.step()
神经网络验证回路
# 将模型设置为评估模式
model.eval()
total_loss = 0
# 通过数据加载器迭代
for x, y in dv_set:
# cpu/cuda(gpu)
x, y = x.to(device), y.to(device)
# 禁用梯度计算
with torch.no_grad():
# 前向传播(计算输出)
pred = model(x)
# 计算损失
loss = criterion(pred, y)
# 累积损失
total_loss += loss.cpu().item() * len(x)
# 计算平均损失
avg_loss = total_loss / len(dv_set.dataset)
神经网络测试回路
# # 将模型设置为评估模式
model.eval()
preds = []
for x in tt_set:
x = x.to(device)
with torch.no_grad():
pred = model(x)
# 收集预测结果
preds.append(pred.cpu())