wordpress响应式concise主题,做网站优化时代码结构关系大吗,西安微信商城网站开发,soho建设外贸网站之前学习了RadioDiff这篇论文#xff0c;最近在复刻相关代码。
这段代码实现了一个基于潜在扩散模型#xff08;Latent Diffusion Model, LDM#xff09;的训练框架。借助DeepSeek总体学习一下#xff1a;
1. 整体结构
代码主要分为以下几个部分#xff1a;
参数解析和…之前学习了RadioDiff这篇论文最近在复刻相关代码。
这段代码实现了一个基于潜在扩散模型Latent Diffusion Model, LDM的训练框架。借助DeepSeek总体学习一下
1. 整体结构
代码主要分为以下几个部分
参数解析和配置加载parse_args, load_conf主训练函数main训练器类Trainer主程序入口
2. 主要组件
2.1 模型架构
AutoencoderKL变分自编码器用于将图像编码到潜在空间Unet条件UNet模型用于扩散过程LatentDiffusion将VAE和UNet组合在一起的潜在扩散模型
2.2 训练流程
加载配置文件和模型准备数据集和数据加载器初始化训练器执行训练循环
3. 详细解释
3.1 参数和配置
def parse_args():parser argparse.ArgumentParser(descriptiontraining vae configure)parser.add_argument(--cfg, helpexperiment configure file name, typestr, requiredTrue)args parser.parse_args()args.cfg load_conf(args.cfg)return args使用argparse解析命令行参数主要接受一个配置文件路径load_conf函数加载YAML配置文件
3.2 主函数
def main(args):cfg CfgNode(args.cfg)# 初始化VAE模型first_stage_model AutoencoderKL(...)# 初始化UNet模型if model_cfg.model_name cond_unet:from denoising_diffusion_pytorch.mask_cond_unet import Unetunet Unet(...)# 初始化扩散模型if model_cfg.model_type const_sde:from denoising_diffusion_pytorch.ddm_const_sde import LatentDiffusionldm LatentDiffusion(...)# 统计模型参数vae_params count_params(first_stage_model)unet_params count_params(unet)ldm_total_params count_params(ldm)# 准备数据集和数据加载器if data_cfg[name] edge:dataset EdgeDataset(...)dl DataLoader(...)# 初始化训练器并开始训练trainer Trainer(...)trainer.train()3.3 Trainer类
Trainer类负责整个训练过程的管理主要功能包括 初始化 设置混合精度训练准备模型、数据加载器、优化器初始化EMA指数移动平均模型设置学习率调度器 训练循环 def train(self):while self.step self.train_num_steps:# 梯度累积for ga_ind in range(self.gradient_accumulate_every):batch next(self.dl)# 训练步骤loss, log_dict self.model.training_step(batch)# 反向传播self.accelerator.backward(loss)# 参数更新self.opt.step()self.opt.zero_grad()self.lr_scheduler.step()# EMA更新self.ema.update()# 定期保存和采样if self.step % self.save_and_sample_every 0:self.save(milestone)# 生成样本图像all_images, *_ self.model.sample(...)tv.utils.save_image(all_images, ...)关键功能 梯度累积通过多次前向传播累积梯度再更新参数混合精度训练使用accelerator自动管理EMA保持模型的平滑版本提高稳定性定期采样生成样本图像监控训练进度
3.4 模型保存和加载
def save(self, milestone):data {step: self.step,model: self.accelerator.get_state_dict(self.model),opt: self.opt.state_dict(),ema: self.ema.state_dict(),...}torch.save(data, fmodel-{milestone}.pt)def load(self, milestone):data torch.load(...)model.load_state_dict(data[model])self.opt.load_state_dict(data[opt])...4. 关键技术点 潜在扩散模型 先在潜在空间进行扩散过程再通过VAE解码降低计算成本提高生成质量 条件生成 UNet支持条件输入cond_in_dim, cond_dim等参数可用于图像到图像转换任务 分布式训练 使用Accelerator库简化多GPU训练支持混合精度训练 训练技巧 EMA模型平滑学习率调度梯度裁剪
5. 配置文件
代码通过YAML文件配置模型和训练参数包括
模型架构参数训练超参数数据路径和预处理设置
6. 应用场景
这段代码适用于
图像生成图像修复图像到图像转换其他基于扩散模型的生成任务
总结来说这是一个完整的潜在扩散模型训练框架包含了从数据加载、模型定义、训练循环到结果保存的全流程实现支持分布式训练和各种训练优化技术。
完整代码
import yaml
import argparse
import math
import torch
from lib import loaders
import torch.nn as nn
from tqdm.auto import tqdm
from denoising_diffusion_pytorch.ema import EMA
from accelerate import Accelerator, DistributedDataParallelKwargs
from torch.utils.tensorboard import SummaryWriter
from denoising_diffusion_pytorch.utils import *
import torchvision as tv
from denoising_diffusion_pytorch.encoder_decoder import AutoencoderKL
# from denoising_diffusion_pytorch.transmodel import TransModel
from denoising_diffusion_pytorch.data import *
from torch.utils.data import DataLoader
from multiprocessing import cpu_count
from fvcore.common.config import CfgNode
from taming.modules.util import count_params
os.environ[CUDA_VISIBLE_DEVICES]1def parse_args():parser argparse.ArgumentParser(descriptiontraining vae configure)parser.add_argument(--cfg, helpexperiment configure file name, typestr, requiredTrue)# parser.add_argument()args parser.parse_args()args.cfg load_conf(args.cfg)return argsdef load_conf(config_file, conf{}):with open(config_file) as f:exp_conf yaml.load(f, Loaderyaml.FullLoader)for k, v in exp_conf.items():conf[k] vreturn confdef main(args):cfg CfgNode(args.cfg)# logger create_logger(root_dircfg[out_path])# writer SummaryWriter(cfg[out_path])model_cfg cfg.modelfirst_stage_cfg model_cfg.first_stagefirst_stage_model AutoencoderKL(ddconfigfirst_stage_cfg.ddconfig,lossconfigfirst_stage_cfg.lossconfig,embed_dimfirst_stage_cfg.embed_dim,ckpt_pathfirst_stage_cfg.ckpt_path,)if model_cfg.model_name cond_unet:from denoising_diffusion_pytorch.mask_cond_unet import Unetunet_cfg model_cfg.unetunet Unet(dimunet_cfg.dim,channelsunet_cfg.channels,dim_multsunet_cfg.dim_mults,learned_varianceunet_cfg.get(learned_variance, False),out_mulunet_cfg.out_mul,cond_in_dimunet_cfg.cond_in_dim,cond_dimunet_cfg.cond_dim,cond_dim_multsunet_cfg.cond_dim_mults,window_sizes1unet_cfg.window_sizes1,window_sizes2unet_cfg.window_sizes2,fourier_scaleunet_cfg.fourier_scale,cfgunet_cfg,)else:raise NotImplementedErrorif model_cfg.model_type const_sde:from denoising_diffusion_pytorch.ddm_const_sde import LatentDiffusionelse:raise NotImplementedError(f{model_cfg.model_type} is not surportted !)ldm LatentDiffusion(modelunet,auto_encoderfirst_stage_model,train_samplemodel_cfg.train_sample,image_sizemodel_cfg.image_size,timestepsmodel_cfg.timesteps,sampling_timestepsmodel_cfg.sampling_timesteps,loss_typemodel_cfg.loss_type,objectivemodel_cfg.objective,scale_factormodel_cfg.scale_factor,scale_by_stdmodel_cfg.scale_by_std,scale_by_softsignmodel_cfg.scale_by_softsign,default_scalemodel_cfg.get(default_scale, False),input_keysmodel_cfg.input_keys,ckpt_pathmodel_cfg.ckpt_path,ignore_keysmodel_cfg.ignore_keys,only_modelmodel_cfg.only_model,start_distmodel_cfg.start_dist,perceptual_weightmodel_cfg.perceptual_weight,use_l1model_cfg.get(use_l1, True),cfgmodel_cfg,)# 统计VAEAutoencoderKLvae_params count_params(first_stage_model)print(f扩散模型中的VAE参数量: {vae_params / 1e6:.2f}M)# 统计UNetunet_params count_params(unet)print(fUNet参数量: {unet_params:,})print(fUNet参数量: {unet_params / 1e6:.2f}M)# 统计整个扩散模型包括UNet、VAE等ldm_total_params count_params(ldm)print(f扩散模型总参数量: {ldm_total_params:,})print(f扩散模型总参数量: {ldm_total_params / 1e6:.2f}M)data_cfg cfg.dataif data_cfg[name] edge:dataset EdgeDataset(data_rootdata_cfg.img_folder,image_sizemodel_cfg.image_size,augment_horizontal_flipdata_cfg.augment_horizontal_flip,cfgdata_cfg)elif data_cfg[name] radio:dataset loaders.RadioUNet_c(phasetrain)else:raise NotImplementedErrordl DataLoader(dataset, batch_sizedata_cfg.batch_size, shuffleTrue, pin_memoryTrue,num_workersdata_cfg.get(num_workers, 2))train_cfg cfg.trainertrainer Trainer(ldm, dl, train_batch_sizedata_cfg.batch_size,gradient_accumulate_everytrain_cfg.gradient_accumulate_every,train_lrtrain_cfg.lr, train_num_stepstrain_cfg.train_num_steps,save_and_sample_everytrain_cfg.save_and_sample_every, results_foldertrain_cfg.results_folder,amptrain_cfg.amp, fp16train_cfg.fp16, log_freqtrain_cfg.log_freq, cfgcfg,resume_milestonetrain_cfg.resume_milestone,train_wdtrain_cfg.get(weight_decay, 1e-4))if train_cfg.test_before:if trainer.accelerator.is_main_process:with torch.no_grad():for datatmp in dl:breakif isinstance(trainer.model, nn.parallel.DistributedDataParallel):all_images, *_ trainer.model.module.sample(batch_sizedatatmp[cond].shape[0],conddatatmp[cond].to(trainer.accelerator.device),maskdatatmp[ori_mask].to(trainer.accelerator.device) if ori_mask in datatmp else None)elif isinstance(trainer.model, nn.Module):all_images, *_ trainer.model.sample(batch_sizedatatmp[cond].shape[0],conddatatmp[cond].to(trainer.accelerator.device),maskdatatmp[ori_mask].to(trainer.accelerator.device) if ori_mask in datatmp else None)# all_images torch.cat(all_images_list, dim 0)nrow 2 ** math.floor(math.log2(math.sqrt(data_cfg.batch_size)))tv.utils.save_image(all_images, str(trainer.results_folder / fsample-{train_cfg.resume_milestone}_{model_cfg.sampling_timesteps}.png), nrownrow)torch.cuda.empty_cache()trainer.train()passclass Trainer(object):def __init__(self,model,data_loader,train_batch_size16,gradient_accumulate_every1,train_lr1e-4,train_wd1e-4,train_num_steps100000,save_and_sample_every1000,num_samples25,results_folder./results,ampFalse,fp16False,split_batchesTrue,log_freq20,resume_milestone0,cfg{},):super().__init__()ddp_handler DistributedDataParallelKwargs(find_unused_parametersTrue)self.accelerator Accelerator(split_batchessplit_batches,mixed_precisionfp16 if fp16 else no,kwargs_handlers[ddp_handler],)self.enable_resume cfg.trainer.get(enable_resume, False)self.accelerator.native_amp ampself.model modelassert has_int_squareroot(num_samples), number of samples must have an integer square rootself.num_samples num_samplesself.save_and_sample_every save_and_sample_everyself.batch_size train_batch_sizeself.gradient_accumulate_every gradient_accumulate_everyself.log_freq log_freqself.train_num_steps train_num_stepsself.image_size model.image_size# dataset and dataloader# self.ds Dataset(folder, mask_folder, self.image_size, augment_horizontal_flip augment_horizontal_flip, convert_image_to convert_image_to)# dl DataLoader(self.ds, batch_size train_batch_size, shuffle True, pin_memory True, num_workers cpu_count())dl self.accelerator.prepare(data_loader)self.dl cycle(dl)# optimizerself.opt torch.optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()),lrtrain_lr, weight_decaytrain_wd)lr_lambda lambda iter: max((1 - iter / train_num_steps) ** 0.96, cfg.trainer.min_lr)self.lr_scheduler torch.optim.lr_scheduler.LambdaLR(self.opt, lr_lambdalr_lambda)# for logging results in a folder periodicallyif self.accelerator.is_main_process:self.results_folder Path(results_folder)self.results_folder.mkdir(exist_okTrue, parentsTrue)self.ema EMA(model, ema_modelNone, beta0.999,update_after_stepcfg.trainer.ema_update_after_step,update_everycfg.trainer.ema_update_every)# step counter stateself.step 0# prepare model, dataloader, optimizer with acceleratorself.model, self.opt, self.lr_scheduler \self.accelerator.prepare(self.model, self.opt, self.lr_scheduler)self.logger create_logger(root_dirresults_folder)self.logger.info(cfg)self.writer SummaryWriter(results_folder)self.results_folder Path(results_folder)resume_file str(self.results_folder / fmodel-{resume_milestone}.pt)if os.path.isfile(resume_file):self.load(resume_milestone)def save(self, milestone):if not self.accelerator.is_local_main_process:returnif self.enable_resume:data {step: self.step,model: self.accelerator.get_state_dict(self.model),opt: self.opt.state_dict(),lr_scheduler: self.lr_scheduler.state_dict(),ema: self.ema.state_dict(),scaler: self.accelerator.scaler.state_dict() if exists(self.accelerator.scaler) else None}# data_only_model {ema: self.ema.state_dict(),}torch.save(data, str(self.results_folder / fmodel-{milestone}.pt))else:data {model: self.accelerator.get_state_dict(self.model),}torch.save(data, str(self.results_folder / fmodel-{milestone}.pt))def load(self, milestone):assert self.enable_resume; resume is available only if self.enable_resume is True !accelerator self.acceleratordevice accelerator.devicedata torch.load(str(self.results_folder / fmodel-{milestone}.pt),map_locationlambda storage, loc: storage)model self.accelerator.unwrap_model(self.model)model.load_state_dict(data[model])if scale_factor in data[model]:model.scale_factor data[model][scale_factor]self.step data[step]self.opt.load_state_dict(data[opt])self.lr_scheduler.load_state_dict(data[lr_scheduler])if self.accelerator.is_main_process:self.ema.load_state_dict(data[ema])if exists(self.accelerator.scaler) and exists(data[scaler]):self.accelerator.scaler.load_state_dict(data[scaler])def train(self):accelerator self.acceleratordevice accelerator.devicewith tqdm(initialself.step, totalself.train_num_steps, disablenot accelerator.is_main_process) as pbar:while self.step self.train_num_steps:total_loss 0.total_loss_dict {loss_simple: 0., loss_vlb: 0., total_loss: 0., lr: 5e-5}for ga_ind in range(self.gradient_accumulate_every):# data next(self.dl).to(device)batch next(self.dl)for key in batch.keys():if isinstance(batch[key], torch.Tensor):batch[key].to(device)if self.step 0 and ga_ind 0:if isinstance(self.model, nn.parallel.DistributedDataParallel):self.model.module.on_train_batch_start(batch)else:self.model.on_train_batch_start(batch)with self.accelerator.autocast():if isinstance(self.model, nn.parallel.DistributedDataParallel):loss, log_dict self.model.module.training_step(batch)else:loss, log_dict self.model.training_step(batch)loss loss / self.gradient_accumulate_everytotal_loss loss.item()loss_simple log_dict[train/loss_simple].item() / self.gradient_accumulate_everyloss_vlb log_dict[train/loss_vlb].item() / self.gradient_accumulate_everytotal_loss_dict[loss_simple] loss_simpletotal_loss_dict[loss_vlb] loss_vlbtotal_loss_dict[total_loss] total_loss# total_loss_dict[s_fact] self.model.module.scale_factor# total_loss_dict[s_bias] self.model.module.scale_biasself.accelerator.backward(loss)total_loss_dict[lr] self.opt.param_groups[0][lr]describtions dict2str(total_loss_dict)describtions [Train Step] {}/{}: .format(self.step, self.train_num_steps) describtionsif accelerator.is_main_process:pbar.desc describtionsif self.step % self.log_freq 0:if accelerator.is_main_process:# pbar.desc describtions# self.logger.info(pbar.__str__())self.logger.info(describtions)accelerator.clip_grad_norm_(filter(lambda p: p.requires_grad, self.model.parameters()), 1.0)# pbar.set_description(floss: {total_loss:.4f})accelerator.wait_for_everyone()self.opt.step()self.opt.zero_grad()self.lr_scheduler.step()if accelerator.is_main_process:self.writer.add_scalar(Learning_Rate, self.opt.param_groups[0][lr], self.step)self.writer.add_scalar(total_loss, total_loss, self.step)self.writer.add_scalar(loss_simple, loss_simple, self.step)self.writer.add_scalar(loss_vlb, loss_vlb, self.step)accelerator.wait_for_everyone()self.step 1# if self.step int(self.train_num_steps * 0.2):if accelerator.is_main_process:self.ema.to(device)self.ema.update()if self.step ! 0 and self.step % self.save_and_sample_every 0:milestone self.step // self.save_and_sample_everyself.save(milestone)self.model.eval()# self.ema.ema_model.eval()with torch.no_grad():# img self.dl# batches num_to_groups(self.num_samples, self.batch_size)# all_images_list list(map(lambda n: self.model.module.validate_img(nsself.batch_size), batches))if isinstance(self.model, nn.parallel.DistributedDataParallel):# all_images self.model.module.sample(batch_sizeself.batch_size)all_images, *_ self.model.module.sample(batch_sizebatch[cond].shape[0],condbatch[cond],maskbatch[ori_mask] if ori_mask in batch else None)elif isinstance(self.model, nn.Module):# all_images self.model.sample(batch_sizeself.batch_size)all_images, *_ self.model.sample(batch_sizebatch[cond].shape[0],condbatch[cond],maskbatch[ori_mask] if ori_mask in batch else None)# all_images torch.clamp((all_images 1.0) / 2.0, min0.0, max1.0)# all_images torch.cat(all_images_list, dim 0)# nrow 2 ** math.floor(math.log2(math.sqrt(self.batch_size)))nrow 2 ** math.floor(math.log2(math.sqrt(batch[cond].shape[0])))tv.utils.save_image(all_images, str(self.results_folder / fsample-{milestone}.png), nrownrow)self.model.train()accelerator.wait_for_everyone()pbar.update(1)accelerator.print(training complete)if __name__ __main__:args parse_args()main(args)pass
