Skip to content

Latest commit

 

History

History
75 lines (54 loc) · 2.26 KB

调试日志.md

File metadata and controls

75 lines (54 loc) · 2.26 KB

criterion_prob = nn.BCELoss() # 二元交叉熵损失用于概率 criterion_time = nn.MSELoss() # 均方误差损失用于时间

检查是否有可用的GPU

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f'使用设备: {device}')

将损失函数移动到GPU

criterion_prob = criterion_prob.to(device) criterion_time = criterion_time.to(device)

将优化器中的参数移动到GPU

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

print(f'模型、损失函数和优化器已移动到 {device}')

将模型移动到GPU

model = model.to(device)

修改训练循环以使用GPU

num_epochs = 10 for epoch in range(num_epochs): for g in train_dataloader: g = g.to(device) # 将图移动到GPU g = dgl.add_self_loop(g) probs, times = model(g, g.ndata["feat"].float()) # 将特征移动到GPU

    t = g.ndata["label"][0]  # 将标签移动到GPU
    
    true_probs = t[::2]
    true_times = t[1::2]
    
    probs = probs.squeeze(0)
    times = times.squeeze(0)
    
    loss_prob = criterion_prob(probs, true_probs)
    loss_time = criterion_time(times, true_times)
    loss = loss_prob + loss_time
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
print(f"第 {epoch} 轮: 损失 = {loss.item()}")

在训练结束后,如果需要在CPU上使用模型,可以将其移回CPU

model = model.to('cpu')

from tabulate import tabulate

for epoch in range(start_epoch, num_epochs): for g in train_dataloader: # ... 前面的代码保持不变 ...

    probs = probs.squeeze(0)
    times = times.squeeze(0)
    
    # 准备表格数据
    table_data = []
    for i in range(len(true_probs)):
        table_data.append([
            f"策略 {i+1}",
            f"{true_probs[i]:.4f}",
            f"{probs[i].item():.4f}",
            f"{true_times[i]:.4f}",
            f"{times[i].item():.4f}"
        ])
    
    # 打印表格
    headers = ["策略", "真实概率", "预测概率", "真实时间", "预测时间"]
    print(tabulate(table_data, headers=headers, tablefmt="grid"))
    
    # ... 后面的代码保持不变 ...