Deep3DFaceRecon_pytorch-master项目学习-train.py

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()
  • 0
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值