olivame
f86ab51a04
- GroupNorm/LayerNorm bypass autocast，消除 bf16→fp32→bf16 转换开销 - DDIM 调度系数 cast 到输入 dtype，attention mask 直接用 bf16 分配 - alphas_cumprod 提升到 float64 保证数值精度 - SinusoidalPosEmb 输出 dtype跟随模型精度 - 新增 profile_unet.py 脚本及FLOPS 分析结果 - 启用 TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL - case1 PSNR: 30.45 → 30.24（bf16 精度预期内波动）
111 lines
3.1 KiB
Python
111 lines
3.1 KiB
Python
# adopted from
|
|
# https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
|
|
# and
|
|
# https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
|
|
# and
|
|
# https://github.com/openai/guided-diffusion/blob/0ba878e517b276c45d1195eb29f6f5f72659a05b/guided_diffusion/nn.py
|
|
#
|
|
# thanks!
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from unifolm_wma.utils.utils import instantiate_from_config
|
|
|
|
|
|
def disabled_train(self, mode=True):
|
|
"""Overwrite model.train with this function to make sure train/eval mode
|
|
does not change anymore."""
|
|
return self
|
|
|
|
|
|
def zero_module(module):
|
|
"""
|
|
Zero out the parameters of a module and return it.
|
|
"""
|
|
for p in module.parameters():
|
|
p.detach().zero_()
|
|
return module
|
|
|
|
|
|
def scale_module(module, scale):
|
|
"""
|
|
Scale the parameters of a module and return it.
|
|
"""
|
|
for p in module.parameters():
|
|
p.detach().mul_(scale)
|
|
return module
|
|
|
|
|
|
def conv_nd(dims, *args, **kwargs):
|
|
"""
|
|
Create a 1D, 2D, or 3D convolution module.
|
|
"""
|
|
if dims == 1:
|
|
return nn.Conv1d(*args, **kwargs)
|
|
elif dims == 2:
|
|
return nn.Conv2d(*args, **kwargs)
|
|
elif dims == 3:
|
|
return nn.Conv3d(*args, **kwargs)
|
|
raise ValueError(f"unsupported dimensions: {dims}")
|
|
|
|
|
|
def linear(*args, **kwargs):
|
|
"""
|
|
Create a linear module.
|
|
"""
|
|
return nn.Linear(*args, **kwargs)
|
|
|
|
|
|
def avg_pool_nd(dims, *args, **kwargs):
|
|
"""
|
|
Create a 1D, 2D, or 3D average pooling module.
|
|
"""
|
|
if dims == 1:
|
|
return nn.AvgPool1d(*args, **kwargs)
|
|
elif dims == 2:
|
|
return nn.AvgPool2d(*args, **kwargs)
|
|
elif dims == 3:
|
|
return nn.AvgPool3d(*args, **kwargs)
|
|
raise ValueError(f"unsupported dimensions: {dims}")
|
|
|
|
|
|
def nonlinearity(type='silu'):
|
|
if type == 'silu':
|
|
return nn.SiLU()
|
|
elif type == 'leaky_relu':
|
|
return nn.LeakyReLU()
|
|
|
|
|
|
class GroupNormSpecific(nn.GroupNorm):
|
|
|
|
def forward(self, x):
|
|
with torch.amp.autocast('cuda', enabled=False):
|
|
return F.group_norm(x, self.num_groups,
|
|
self.weight.to(x.dtype) if self.weight is not None else None,
|
|
self.bias.to(x.dtype) if self.bias is not None else None,
|
|
self.eps)
|
|
|
|
|
|
def normalization(channels, num_groups=32):
|
|
"""
|
|
Make a standard normalization layer.
|
|
:param channels: number of input channels.
|
|
:return: an nn.Module for normalization.
|
|
"""
|
|
return GroupNormSpecific(num_groups, channels)
|
|
|
|
|
|
class HybridConditioner(nn.Module):
|
|
|
|
def __init__(self, c_concat_config, c_crossattn_config):
|
|
super().__init__()
|
|
self.concat_conditioner = instantiate_from_config(c_concat_config)
|
|
self.crossattn_conditioner = instantiate_from_config(
|
|
c_crossattn_config)
|
|
|
|
def forward(self, c_concat, c_crossattn):
|
|
c_concat = self.concat_conditioner(c_concat)
|
|
c_crossattn = self.crossattn_conditioner(c_crossattn)
|
|
return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]}
|