pycharm进度条使用
一、pytorch 深度学习训练如何显示进度条
1、通过使用tadm，实时显示训练进度，并显示当前训练集正确率以及损失等信息
大概就是这么一个效果
要导入这个包
from tqdm import tqdm
一般就是在epoch那个for循环里面写
2、整个网络代码（仅作定位参考，比如在哪里写啊，在那里更新啊，具体看看这个代码，然后根据自己的代码适当修改）
import torch
import torch.nn as nn
import os
import numpy as np
from models import spinal_net
import decoder
import loss
from dataset import BaseDataset
from tqdm import tqdm
from torch.utils.tensorboard import SummaryWriter # 查看中间的输出过程
import time
# ---------------------------------------------------------------------------------------------------
def collater(data):
"""
将一个batch的数据进行整合，返回一个字典，包含了每个数据样本中对应名字的所有数据，并且每个数据的第一维度都是batch size。
"""
out_data_dict = {}
# # 创建一个包含所有数据名称的字典，值为一个空列表
for name in data[0]:
out_data_dict[name] = []
# # 将每个数据样本的数据按照名称，添加到字典的对应名称下
for sample in data:
for name in sample:
out_data_dict[name].append(torch.from_numpy(sample[name]))
# # 将字典中每个名称对应的数据列表，沿着第0维度拼接成一个tensor
for name in out_data_dict:
out_data_dict[name] = torch.stack(out_data_dict[name], dim=0)
# # 返回整合后的字典
return out_data_dict
# ------------------------------------------------------------------------------------------
# 创建网络模型
# ------------------------------------------------------------------------------------------
class Network(object):
# 初始化网络
def __init__(self, args):
torch.manual_seed(317) # 设置PyTorch中的随机数种子，以便在每次运行代码时生成相同的随机数序列。
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 这个head定义的是一个的
heads = {'hm': args.num_classes, # 68个关键点，68class,热图
'reg': 2*args.num_classes, # 回归的中心点x,y偏移？
'wh': 2*4,} # 四个点的回归x，y偏移？
# 主要用的就是这个网络
self.model = spinal_net.SpineNet(heads=heads,
pretrained=True,
down_ratio=args.down_ratio,
final_kernel=1,
head_conv=256)
self.num_classes = args.num_classes
self.decoder = decoder.DecDecoder(K=args.K, conf_thresh=args.conf_thresh) # 这个是什么？ conf_thresh认值为 0.2，用于设置置信度阈值。
# 将网络的输出进行解码成我们想要的，上述代码就是初始化了一下
self.dataset = {'spinal': BaseDataset} # 可以进去看看怎么制作数据的
# ----------------------------------------------------------------
# 保存权重
def save_model(self, path, epoch, model):
"""
这个函数的作用是保存模型的权重到指定的路径。它接受三个参数：
path：保存模型权重的路径
epoch：当前的训练轮数
model：需要保存权重的模型
"""
if isinstance(model, torch.nn.DataParallel):
state_dict = model.module.state_dict()
else:
state_dict = model.state_dict()
data = {'epoch': epoch, 'state_dict': state_dict}
torch.save(data, path)
# -------------------------------------------------------------------------------
# 加载模型
def load_model(self, model, resume, strict=True):
"""
主要用于加载预训练模型的权重到现有模型中
如果strict是True，那么将只加载与现有模型完全匹配的参数，
否则可以使用预训练模型中不匹配的参数。该函数首先读取预训练模型的权重，
然后通过将现有模型的参数名从字符串“module”中删除前缀“module”来对其进行调整。
接下来，它将根据strict参数加载权重，并将结果返回。
"""
checkpoint = torch.load(resume, map_location=lambda storage, loc: storage)
# 这行代码的作用是从指定路径resume中加载一个已保存的 PyTorch 模型的参数。
# map_location参数允许我们指定设备类型，以便将模型参数映射到正确的设备上。如果模型在 GPU 上训练，
# 但是当前没有 GPU 可用，我们可以使用 map_location=torch.device('cpu') 将参数加载到 CPU 上。
print('loaded weights from {}, epoch {}'.format(resume, checkpoint['epoch']))
state_dict_ = checkpoint['state_dict']
state_dict = {}
for k in state_dict_:
if k.startswith('module') and not k.startswith('module_list'):
state_dict[k[7:]] = state_dict_[k]
else:
state_dict[k] = state_dict_[k]
model_state_dict = model.state_dict()
if not strict:
for k in state_dict:
if k in model_state_dict:
if state_dict[k].shape != model_state_dict[k].shape:
print('Skip loading parameter {}, required shape{}, ' \
'loaded shape{}.'.format(k, model_state_dict[k].shape, state_dict[k].shape))
state_dict[k] = model_state_dict[k]
else:
print('Drop parameter {}.'.format(k))
for k in model_state_dict:
if not (k in state_dict):
print('No param {}.'.format(k))
state_dict[k] = model_state_dict[k]
model.load_state_dict(state_dict, strict=False)
return model
#------------------------------------------- 开始训练网络了 ------------------------------------------------
def train_network(self, args): # main()进入到这里
device = self.device
print("using {} device.".format(device))
# ----------------------------TensorBoard(dxp)-------------------------------------------
project_path = r"F:\py_project_professional\dxp_project\Vertebra-Landmark-Detection-master"
name = "Vertebra-Landmark-Detection-master"
tb_writer_summary_path = os.path.join(project_path, "run", name, "Logs")
current_time = time.strftime("%Y%m%d-%H%M%S", time.localtime())
log_dir = os.path.join(tb_writer_summary_path, current_time)
tb_writer = SummaryWriter(log_dir=log_dir, comment=name)
# # 于将训练过程中的各种信息写入 TensorBoard 日志文件，以便进行可视化和分析。
# # log_dir：用于指定 TensorBoard 日志文件的保存路径，即 tb_writer_summary_path 变量所指定的路径。
# # comment：一个可选的字符串，用于为日志文件添加一个注释，可以用于区分不同的日志文件。
# # 在这里，该参数指定了一个名为 args.name 的注释，该值可能是从命令行参数中传递过来的。
# # parser.add_argument('--name', type=str, default='case5_result')
# ---------------------------------------------------------------------------------------
save_path = 'weights_'+args.dataset # 保存权重的路径(weight_spinal)
if not os.path.exists(save_path): # 没有就创建这个文件夹
os.mkdir(save_path)
self.optimizer = torch.optim.Adam(self.model.parameters(), args.init_lr)
scheduler = torch.optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=0.96, last_epoch=-1)
if args.ngpus>0: # 一般不执行
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
self.model = nn.DataParallel(self.model)
self.model.to(self.device) # 模型指认到设备
# -----------------------------------------
# print(self.model) # 输出一下这个模型看看？
# ----------------------------------------
criterion = loss.LossAll() # 损失（这个损失是自己写的），这只是初始化一下损失函数
# 这个损失函数进去看看
print('Setting up data...（开始加载数据）') # 能跑到.......
dataset_module = self.dataset[args.dataset]
# -------------------训练开始怎么输出train数据有多少张，val数据有多少张---(dxp)
dsets = {
'train': dataset_module(data_dir=args.data_dir, phase='train', input_h=args.input_h, input_w=args.input_w,
down_ratio=args.down_ratio),
'val': dataset_module(data_dir=args.data_dir, phase='val', input_h=args.input_h, input_w=args.input_w,
down_ratio=args.down_ratio)
}
train_num = len(dsets['train'])
val_num = len(dsets['val'])
print('using {} images for training'.format(train_num))
print('using {} images for validation'.format(val_num) )
# ----------------------------------------------------------------
# 进入数据的处理，制作，读文件，读关键点，然后预处理，然后生成GT
dsets = {x: dataset_module(data_dir=args.data_dir,
phase=x,
input_h=args.input_h,
input_w=args.input_w,
down_ratio=args.down_ratio)
for x in ['train', 'val']}
# 进入了这个函数class BaseDataset(data.Dataset):
# 这段代码定义了一个字典dsets，其中包含两个键值对，分别对应训练集和验证集。
# 对于每一个数据集，使用了dataset_module函数，传入了一些参数，包括数据集的存储目录、数据预处理后的图像尺寸、下采样率等等，返回的是一个数据集实例。
# 数据记载器，这也是初始化
dsets_loader = {'train': torch.utils.data.DataLoader(dsets['train'],
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True, # 如果在时期结束时存在不完整的批次，则丢弃这些批次
collate_fn=collater), #这是一个自定义函数，用于将每个批次中的样本汇总成一个批次张量。
'val':torch.utils.data.DataLoader(dsets['val'],
batch_size=1,
shuffle=False,
num_workers=1,
pin_memory=True,
collate_fn=collater)}
print('Starting training...（开始训练）')
train_loss = []
val_loss = []
# 真正开始跑
for epoch in range(1, args.num_epoch+1):
print('-'*75)
print('Epoch: {}/{} '.format(epoch, args.num_epoch)) # # 又能跑到这......................
# ----------------------------------------------------------------------------------
# 使用tqdm创建进度条（dxp）
dsets_loader['train'] = tqdm(dsets_loader['train'], total=train_num,
desc=f"train epoch[{epoch}/{args.num_epoch}]")
# -----------------------进入run_epoch函数,训练集的epoch-------------------------------
tr_epoch_loss = self.run_epoch(phase='train',
data_loader=dsets_loader['train'],
criterion=criterion)
# 返回的是一个epoch_loss
# 进去run_epoch 函数进行迭代训练(训练一轮后再出来)
train_loss.append(tr_epoch_loss)
# scheduler.step(epoch)
scheduler.step()
# ---------------------------------------------------------------------------------
# 使用tqdm创建进度条（dxp）
dsets_loader['val'] = tqdm(dsets_loader['val'], total=val_num, desc=f"val epoch[{epoch}/{args.num_epoch}]")
# -------------------------------测试集的epoch---------------------------------------
val_epoch_loss = self.run_epoch(phase='val',
data_loader=dsets_loader['val'],
criterion=criterion)
val_loss.append(val_epoch_loss)
# ----------------------------------------------------------------------------------
# 将训练和验证的损失损失保存在
np.savetxt(os.path.join(save_path, 'train_loss.txt'), train_loss, fmt='%.6f')
np.savetxt(os.path.join(save_path, 'val_loss.txt'), val_loss, fmt='%.6f')
# --------------(dxp 写 tensorboard)
tb_writer.add_scalar("train_loss", tr_epoch_loss, epoch)
tb_writer.add_scalar("train_loss", val_epoch_loss, epoch)
# ---------------------------
if epoch % 10 == 0 or epoch ==1:
self.save_model(os.path.join(save_path, 'model_{}.pth'.format(epoch)), epoch, self.model)
if len(val_loss)>1:
if val_loss[-1]<np.min(val_loss[:-1]):
# self.save_model(os.path.join(save_path, 'model_last.pth'), epoch, self.model)
self.save_model(os.path.join(save_path, 'model_best.pth'), epoch, self.model)
# 这段代码的作用是在训练过程中保存验证集上最好的模型，即在验证集上得到最小的loss值的模型。
# 具体实现是在每次完成一轮验证集的训练后，判断当前的验证集loss是否比之前所有轮次中的最小值还要小，如果是，则保存当前模型。
# val_loss是用来记录每轮验证集上的loss值的列表。
# 关闭TensorBoard写入器
tb_writer.close()
# ---------------------------------------------开始跑epoch--------------------------------------------
def run_epoch(self, phase, data_loader, criterion):
if phase == 'train':
self.model.train()
else:
self.model.eval()
running_loss = 0.
# ----------------dxp进度条---------------------------
# 计算总批次数
# total_batches = len(data_loader)
# 使用tqdm创建进度条
# data_loader = tqdm(data_loader, total=total_batches, desc=f"{phase} epoch")
# -------------------------------------------------
# for data_dict in data_loader: # 这里还没出来
for batch_idx, data_dict in enumerate(data_loader): # （dxp）
for name in data_dict:
data_dict[name] = data_dict[name].to(device=self.device)
# 字典的key；input、hm、ind、reg、wh、reg_mask(这是在dataset写好的gititem方法)
if phase == 'train':
self.optimizer.zero_grad()
with torch.enable_grad():
pr_decs = self.model(data_dict['input']) # input是图片，输入到model--->输出pr_decs
loss = criterion(pr_decs, data_dict)
loss.backward()
self.optimizer.step()
else:
with torch.no_grad():
pr_decs = self.model(data_dict['input'])
loss = criterion(pr_decs, data_dict) # 是一个损失函数，用于衡量预测结果与目标数据之间的差异。
running_loss += loss.item()
# 更新进度条的描述（dxp）
data_loader.set_postfix({'loss': loss.item()})
epoch_loss = running_loss / len(data_loader)
print('{} loss: {}'.format(phase, epoch_loss))
return epoch_loss
4168
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
