# 导入必要的库
import torch
import torch.nn as nn
import torch.optim as optim
from time import time
# 模型、损失和优化器
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 设置设备为GPU或CPU
model = DeeperAutoencoderWithMoreAttention().to(device) # 初始化模型并将其移动到设备上
criterion = nn.MSELoss() # 使用均方误差作为损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001) # 使用Adam优化器
best_psnr = 0 # 初始化最佳峰值信噪比（PSNR）
best_model_save_path = './save_root/best_autoencoder.pth' # 最佳模型保存路径
model_save_path = './save_root/epoch{}_model.pth' # 每个epoch的模型保存路径
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=3000, gamma=0.1) # 学习率调整器
num_epochs = 5000 # 总共的训练轮数
if __name__ == '__main__':
print(device)
print("开始运行ing")
for epoch in range(num_epochs): # 循环训练每个epoch
start_time = time() # 记录当前时间
for (clean_images, _), (noisy_images, _) in zip(clean_loader, noisy_loader):
clean_images = clean_images.to(device) # 将干净图像移动到设备上
noisy_images = noisy_images.to(device) # 将含噪图像移动到设备上
outputs = model(noisy_images) # 前向传播
loss = criterion(outputs, clean_images) # 计算损失
optimizer.zero_grad() # 清零梯度
loss.backward() # 反向传播
optimizer.step() # 更新模型参数
scheduler.step() # 更新学习率
save_combined_images(model, noisy_loader, clean_loader, output_path=f"./save_image/index{epoch}epoch.png") # 保存合并后的图像
# 计算平均PSNR
model.eval()
total_psnr = 0
with torch.no_grad():
for (clean_images, _), (noisy_images, _) in zip(clean_loader, noisy_loader):
clean_images = clean_images.to(device)
noisy_images = noisy_images.to(device)
outputs = model(noisy_images)
psnr = compute_psnr(clean_images, outputs)
total_psnr += psnr
avg_psnr = total_psnr / len(clean_loader)
end_time = time()
print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}, PSNR: {avg_psnr:.4f}, Time:{end_time - start_time:.2f}")
if avg_psnr > best_psnr:
best_psnr = avg_psnr
torch.save(model.state_dict(), best_model_save_path)
print(f"在第 {epoch + 1} 个epoch保存了最佳模型，PSNR为: {avg_psnr:.4f}")
if (epoch + 1) % 10 == 0:
torch.save(model.state_dict(), model_save_path.format(epoch))
print(f"在第 {epoch + 1} 个epoch保存模型，PSNR为: {avg_psnr:.4f}")
model.train()
print("Training completed!")
1. 导入必要的库：导入了torch、torch.nn和torch.optim等必要的库。
2. 定义模型、损失和优化器：创建了一个DeeperAutoencoderWithMoreAttention模型，并将其移动到可用的设备上；定义了均方误差作为损失函数；使用Adam优化器来更新模型的参数；初始化了最佳峰值信噪比（PSNR）为0；设置了模型的保存路径；创建了一个学习率调整器scheduler。
3. 训练过程：循环执行num_epochs次训练。
- 对于每个epoch，记录开始时间。
- 遍历干净图像和含噪图像数据加载器，将它们移动到设备上，然后执行前向传播，计算损失并进行反向传播和参数更新。
- 更新学习率。
- 保存合并后的图像。
- 切换模型为评估模式，计算总平均PSNR。
- 打印当前epoch的损失、PSNR和训练时间。
- 如果当前epoch的平均PSNR超过了最佳PSNR，则保存模型，并打印相应信息。
- 如果当前epoch是10的倍数，则保存模型，并打印相应信息。
- 切换模型为训练模式。
4. 训练完成后，输出"Training completed!"。
如果回答对你有帮助，请给我一个最佳答案！
| 
( 粤ICP备18085999号-1 | 粤公网安备 44051102000585号)
狗仔卡
发表于 2023-8-16 19:47:17
置顶卡
千斤顶
显身卡
发表于 2023-8-16 19:49:14