criterion_prob = nn.BCELoss() # 二元交叉熵损失用于概率 criterion_time = nn.MSELoss() # 均方误差损失用于时间
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f'使用设备: {device}')
criterion_prob = criterion_prob.to(device) criterion_time = criterion_time.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
print(f'模型、损失函数和优化器已移动到 {device}')
model = model.to(device)
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()}")
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"))
# ... 后面的代码保持不变 ...