python模块学习
pytorch模块学习
if __name__ == '__main__':
import warnings
warnings.filterwarnings("ignore")
train_opt = TrainOptions().parse() # get training options
world_size = train_opt.world_size
if train_opt.use_ddp:
mp.spawn(main, args=(world_size, train_opt), nprocs=world_size, join=True)
else:
main(0, world_size, train_opt)
def setup(rank, world_size, port):#这个函数是用来设置分布式训练环境的。
os.environ['MASTER_ADDR'] = 'localhost'#这个函数是用来设置分布式训练环境的。
os.environ['MASTER_PORT'] = port#设置主节点端口号为给定的端口号。
dist.init_process_group("gloo", rank=rank, world_size=world_size)#初始化分布式进程组,使用"Gloo"后端,并设置进程的等级和进程组的总大小。
def cleanup():#这个函数用于清除分布式训练环境。
dist.destroy_process_group()#销毁进程组。
# 每个进程都有一个独特的rank值,world_size:这个参数表示分布式训练中进程的总数,train_opt:这个参数可能是一个包含训练选项和超参数的对象,如学习率、批量大小、优化器类型等。
def main(rank, world_size, train_opt):#这是主函数,它将执行训练过程。
val_opt = genvalconf(train_opt, isTrain=False)#生成验证数据集的选项。
#根据给定的等级确定设备。??什么意思
device = torch.device(rank)
#设置当前设备*
torch.cuda.set_device(device)
#检查是否使用分布式数据并行(DDP)
use_ddp = train_opt.use_ddp
if use_ddp:
setup(rank, world_size, train_opt.ddp_port)//设置分布式训练环境。
#创建训练和验证数据集。
train_dataset, val_dataset = create_dataset(train_opt, rank=rank), create_dataset(val_opt, rank=rank)
#分批次
train_dataset_batches, val_dataset_batches = \
len(train_dataset) // train_opt.batch_size, len(val_dataset) // val_opt.batch_size
#根据给定的选项创建模型。
model = create_model(train_opt)
model.setup(train_opt)#设置模型。
model.device = device#设置模型。
model.parallelize()#设置模型。
if rank == 0:
#打印训练和验证图像的批次数量。
print('The batch number of training images = %d\n, \
the batch number of validation images = %d'% (train_dataset_batches, val_dataset_batches))
#打印网络结构
model.print_networks(train_opt.verbose)
#创建一个可视化器,用于显示/保存图像和绘图
visualizer = MyVisualizer(train_opt) # create a visualizer that display/save images and plots
#计算总迭代次数。
total_iters = train_dataset_batches * (train_opt.epoch_count - 1)
t_data = 0#初始化数据加载时间
t_val = 0#初始化验证时间
optimize_time = 0.1#初始化优化时间
batch_size = 1 if train_opt.display_per_batch else train_opt.batch_size#设置批次大小
if use_ddp:
dist.barrier()#添加屏障以同步所有进程。
times = []#初始化空列表记录时间
//始一个循环,遍历每个训练周期(epoch)。
for epoch in range(train_opt.epoch_count, train_opt.n_epochs + 1):
epoch_start_time = time.time() #记录当前时间,用于计算整个训练周期的耗时
iter_data_time = time.time() #记录当前时间,用于计算每次迭代的数据加载时间
epoch_iter = 0 #初始化当前训练周期内的迭代次数,每个周期开始时重置为0
train_dataset.set_epoch(epoch)#初始化当前训练周期内的迭代次数,每个周期开始时重置为
for i, train_data in enumerate(train_dataset): # 开始一个循环,遍历一个训练周期内的所有训练数据。
iter_start_time = time.time() #开始一个循环,遍历一个训练周期内的所有训练数据。
if total_iters % train_opt.print_freq == 0:#判断是否到了打印日志的频率
t_data = iter_start_time - iter_data_time #计算数据加载时间
total_iters += batch_size#更新总迭代次数
epoch_iter += batch_size#更新当前周期的迭代次数
torch.cuda.synchronize()#更新当前周期的迭代次数
optimize_start_time = time.time()#记录当前时间,用于计算优化过程的耗时。
model.set_input(train_data) #记录当前时间,用于计算优化过程的耗时。
model.optimize_parameters() #计算损失函数、获取梯度并更新网络权重
torch.cuda.synchronize()#计算损失函数、获取梯度并更新网络权重
optimize_time = (time.time() - optimize_start_time) / batch_size * 0.005 + 0.995 * optimize_time#计算并更新优化过程的平均耗时。
if use_ddp:
dist.barrier()
#式训练中,只有主进程(即 rank 为 0 的进程)和
#当 total_iters(总迭代次数)等于 batch_size 或者 total_iters 与 train_opt.display_freq(显示频率)的余数为 0 时
if rank == 0 and (total_iters == batch_size or total_iters % train_opt.display_freq == 0):
#调用模型的 compute_visuals 方法来计算当前迭代的可视化结果。这些结果通常包括输入图像、输出图像以及其他可视化信息
model.compute_visuals()
#调用 visualizer(一个 MyVisualizer 实例)的 display_current_results 方法来显示和保存当前迭代的可视化结果。参数说明如下:从模型中获取当前迭代的可视化结果。当前的总迭代次数。当前的训练 epoch。将结果保存到 HTML 文件中。根据 train_opt.add_image 的值决定是否在可视化结果中添加额外的图像。如果为 True,则添加;否则不添加
visualizer.display_current_results(model.get_current_visuals(), total_iters, epoch,
save_results=True,
add_image=train_opt.add_image)
#当total_iters(迄今为止的迭代总数)等于batch_size或total_iters能被train_opt.print_freq整除时
if rank == 0 and (total_iters == batch_size or total_iters % train_opt.print_freq == 0):
#这行代码从当前模型中获取当前的损失值。这通常是一个字典,其中包含了不同损失组件的值。
losses = model.get_current_losses()
#调用visualizer对象的print_current_losses方法来打印当前损失
#参数包括:epoch(当前的训练周期数)、epoch_iter(当前周期的迭代次数)、losses(损失值字典)、optimize_time(优化步骤的耗时)和t_data(可能是数据加载时间)。
visualizer.print_current_losses(epoch, epoch_iter, losses, optimize_time, t_data)
#将当前的损失值绘制成图形。这有助于在训练过程中观察损失的变化趋势。
#参数包括:total_iters(迄今为止的迭代总数)和losses(损失值字典)。
visualizer.plot_current_losses(total_iters, losses)
#这段代码是一个深度学习训练过程中的验证部分,在每个epoch的特定迭代次数上,代码会对验证集进行评估,以获取模型在验证集上的性能。
#1) 当前迭代次数等于批次大小;2) 当前迭代次数能够被评估频率整除。
if total_iters == batch_size or total_iters % train_opt.evaluation_freq == 0:
#在此上下文管理器中,不会计算梯度,这有助于节省内存并提高评估速度。
with torch.no_grad():
torch.cuda.synchronize()#同步所有 CUDA 设备上的操作,确保所有操作都在继续之前完成。
val_start_time = time.time()#记录评估开始时的时间戳
losses_avg = {}#创建一个空字典,用于存储各损失值的平均值
model.eval()#将模型切换到评估模式。
#遍历验证数据集中的每个样本。
for j, val_data in enumerate(val_dataset):
model.set_input(val_data)#将当前验证样本设置为模型的输入。
model.optimize_parameters(isTrain=False)#评估模型性能,不进行优化
#如果当前进程的等级为 0(主进程)且当前批次小于要可视化的批次数,执行以下代码块。
if rank == 0 and j < train_opt.vis_batch_nums:
model.compute_visuals()#计算可视化结果。
#显示当前的可视化结果。
visualizer.display_current_results(model.get_current_visuals(), total_iters, epoch,dataset='val', save_results=True, count=j * val_opt.batch_size,add_image=train_opt.add_image)
#如果当前批次小于要评估的批次数,执行以下代码块。
if j < train_opt.eval_batch_nums:
losses = model.get_current_losses()#获取当前模型损失值
for key, value in losses.items():#遍历遍历损失值字典
losses_avg[key] = losses_avg.get(key, 0) + value#累加每个损失值以计算平均损失。
for key, value in losses_avg.items():#遍历平均损失字典。
#将累加的损失值除以评估的批次数,得到平均损失值。
losses_avg[key] = value / min(train_opt.eval_batch_nums, val_dataset_batches)
torch.cuda.synchronize()#同步所有 CUDA 设备上的操作,确保所有操作都在继续之前完成。
eval_time = time.time() - val_start_time#计算评估所花费的时间
if rank == 0:#如果当前进程的等级为 0(主进程),执行以下代码
#打印当前的损失值
visualizer.print_current_losses(epoch, epoch_iter, losses_avg, eval_time, t_data, dataset='val')
#绘制当前的损失值曲线。
visualizer.plot_current_losses(total_iters, losses_avg, dataset='val')
#将模型设置回训练模
model.train()
#检查是否使用分布式数据并行
if use_ddp:
#如果使用DDP,这里创建一个同步障碍。所有进程必须达到这个障碍,然后才能继续执行。
dist.barrier()
#检查当前进程的等级是否为0(主进程),并且当前迭代次数是否等于批量大小或者是指定保存频率的倍数
if rank == 0 and (total_iters == batch_size or total_iters % train_opt.save_latest_freq == 0):
#检查当前进程的等级是否为0(主进程),并且当前迭代次数是否等于批量大小或者是指定保存频率的倍数
print('saving the latest model (epoch %d, total_iters %d)' % (epoch, total_iters))
#打印实验的名称,便于在控制台跟踪。
print(train_opt.name)
#根据条件选择保存后缀,如果按迭代次数保存,则后缀为迭代次数,否则为'latest'。
save_suffix = 'iter_%d' % total_iters if train_opt.save_by_iter else 'latest'
#保存当前模型
model.save_networks(save_suffix)
#检查是否使用分布式数据并行
if use_ddp:
#如果使用DDP,这里创建一个同步障碍。所有进程必须达到这个障碍,然后才能继续执行。
dist.barrier()
#记录当前时间,用于后续计算耗时。
iter_data_time = time.time()
#打印当前轮次结束的信息,包括当前轮次,总轮次,以及本轮次耗费的时间。
print('End of epoch %d / %d \t Time Taken: %d sec' % (epoch, train_opt.n_epochs, time.time() - epoch_start_time))
#更新学习率
model.update_learning_rate()
#检查当前进程的等级是否为0(主进程),并且当前轮次是否是指定保存轮次频率的倍数。如果满足条件,则执行以下代码块。
if rank == 0 and epoch % train_opt.save_epoch_freq == 0:
#打印保存模型的信息,包括当前轮次(epoch)和迭代次数(total_iters)。
print('saving the model at the end of epoch %d, iters %d' % (epoch, total_iters))
#保存当前模型,后缀为'latest'
model.save_networks('latest')
#保存当前模型,后缀为当前轮次(epoch)
model.save_networks(epoch)
if use_ddp:
dist.barrier()