🚚 (BetaSchedulers) move beta variance schedules in their own module

This commit is contained in:
Laureηt 2023-07-29 15:40:43 +02:00
parent f6ecfdbfda
commit 4dfa3c92dc
Signed by: Laurent
SSH key fingerprint: SHA256:kZEpW8cMJ54PDeCvOhzreNr4FSh6R13CMGH/POoO8DI
10 changed files with 147 additions and 129 deletions

View file

@ -4,7 +4,6 @@
[🤗 diffusers](https://github.com/huggingface/diffusers/), but in Julia (and with worse code) ! [🤗 diffusers](https://github.com/huggingface/diffusers/), but in Julia (and with worse code) !
## Credits ## Credits
- [@huggingface](https://github.com/huggingface) for their amazing [🤗 diffusers](https://github.com/huggingface/diffusers/) library. - [@huggingface](https://github.com/huggingface) for their amazing [🤗 diffusers](https://github.com/huggingface/diffusers/) library.

View file

@ -1,4 +1,5 @@
import Diffusers import Diffusers
import Diffusers.BetaSchedules: cosine_beta_schedule, rescale_zero_terminal_snr
using Flux using Flux
using Random using Random
using Plots using Plots
@ -42,8 +43,8 @@ scatter(data[1, :], data[2, :],
num_timesteps = 100 num_timesteps = 100
scheduler = Diffusers.DDPM( scheduler = Diffusers.DDPM(
Vector{Float64}, Vector{Float64},
Diffusers.rescale_zero_terminal_snr( rescale_zero_terminal_snr(
Diffusers.cosine_beta_schedule(num_timesteps, 0.999f0, 0.001f0) cosine_beta_schedule(num_timesteps, 0.999f0, 0.001f0)
) )
) )

View file

@ -1,125 +0,0 @@
import NNlib: sigmoid
"""
Linear beta schedule.
## Input
* `T::Integer`: number of timesteps
* `β₁::Real=0.0001f0`: initial (t=1) value of β
* `β₋₁::Real=0.02f0`: final (t=T) value of β
## Output
* `β::Vector{Real}`: βₜ values at each timestep t
## References
* [[2006.11239] Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239)
"""
function linear_beta_schedule(T::Integer, β₁::Real=0.0001f0, β₋₁::Real=0.02f0)
return range(start=β₁, stop=β₋₁, length=T)
end
"""
Scaled linear beta schedule.
## Input
* `T::Int`: number of timesteps
* `β₁::Real=0.0001f0`: initial value of β
* `β₋₁::Real=0.02f0`: final value of β
## Output
* `β::Vector{Real}`: βₜ values at each timestep t
## References
* [[2006.11239] Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239)
"""
function scaled_linear_beta_schedule(T::Integer, β₁::Real=0.0001f0, β₋₁::Real=0.02f0)
return range(start=β₁^0.5, stop=β₋₁^0.5, length=T) .^ 2
end
"""
Sigmoid beta schedule.
## Input
* `T::Int`: number of timesteps
* `β₁::Real=0.0001f0`: initial value of β
* `β₋₁::Real=0.02f0`: final value of β
## Output
* `β::Vector{Real}`: βₜ values at each timestep t
## References
* [[2203.02923] GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation](https://arxiv.org/abs/2203.02923)
* [github.com:MinkaiXu/GeoDiff](https://github.com/MinkaiXu/GeoDiff/blob/ea0ca48045a2f7abfccd7f0df449e45eb6eae638/models/epsnet/diffusion.py#L57)
"""
function sigmoid_beta_schedule(T::Integer, β₁::Real=0.0001f0, β₋₁::Real=0.02f0)
x = range(start=-6, stop=6, length=T)
return sigmoid(x) .* (β₋₁ - β₁) .+ β₁
end
"""
Cosine beta schedule.
## Input
* `T::Int`: number of timesteps
* `βₘₐₓ::Real=0.999f0`: maximum value of β
* `ϵ::Real=1e-3f0`: small value used to avoid division by zero
## Output
* `β::Vector{Real}`: βₜ values at each timestep t
## References
* [[2102.09672] Improved Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2102.09672)
* [github:openai/improved-diffusion](https://github.com/openai/improved-diffusion/blob/783b6740edb79fdb7d063250db2c51cc9545dcd1/improved_diffusion/gaussian_diffusion.py#L36)
"""
function cosine_beta_schedule(T::Integer, βₘₐₓ::Real=0.999f0, ϵ::Real=0.001f0)
α̅(t) = cos((t / T + ϵ) / (1 + ϵ) * π / 2)^2
β = Vector{Real}(undef, T)
for t in 1:T
αₜ = α̅(t) / α̅(t - 1)
βₜ = 1 - αₜ
βₜ = min(βₘₐₓ, βₜ)
β[t] = βₜ
end
return β
end
"""
Rescale betas to have zero terminal Signal to Noise Ratio (SNR).
## Input
* `β::AbstractArray`: βₜ values at each timestep t
## Output
* `β::Vector{Real}`: rescaled βₜ values at each timestep t
## References
* [[2305.08891] Rescaling Diffusion Models](https://arxiv.org/abs/2305.08891) (Alg. 1)
"""
function rescale_zero_terminal_snr(β::AbstractArray)
# convert β to ⎷α̅
α = 1 .- β
α̅ = cumprod(α)
⎷α̅ = sqrt.(α̅)
# store old extrema values
⎷α̅₁ = ⎷α̅[1]
⎷α̅₋₁ = ⎷α̅[end]
# shift last timestep to zero
⎷α̅ .-= ⎷α̅₋₁
# scale so that first timestep reaches old values
⎷α̅ *= ⎷α̅₁ / (⎷α̅₁ - ⎷α̅₋₁)
# convert back ⎷α̅ to β
α̅ = ⎷α̅ .^ 2
α = α̅[2:end] ./ α̅[1:end-1]
α = vcat(α̅[1], α)
β = 1 .- α
return β
end

View file

@ -0,0 +1,21 @@
module BetaSchedules
# variance schedulers
include("Linear.jl")
include("ScaledLinear.jl")
include("Cosine.jl")
include("Sigmoid.jl")
export
linear_beta_schedule,
scaled_linear_beta_schedule,
cosine_beta_schedule,
sigmoid_beta_schedule
# utils
include("ZeroSNR.jl")
export
rescale_zero_terminal_snr
end # module BetaSchedules

View file

@ -0,0 +1,30 @@
"""
Cosine beta schedule.
## Input
* `T::Int`: number of timesteps
* `βₘₐₓ::Real=0.999f0`: maximum value of β
* `ϵ::Real=1e-3f0`: small value used to avoid division by zero
## Output
* `β::Vector{Real}`: βₜ values at each timestep t
## References
* [[2102.09672] Improved Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2102.09672)
* [github:openai/improved-diffusion](https://github.com/openai/improved-diffusion/blob/783b6740edb79fdb7d063250db2c51cc9545dcd1/improved_diffusion/gaussian_diffusion.py#L36)
"""
function cosine_beta_schedule(T::Integer, βₘₐₓ::Real=0.999f0, ϵ::Real=0.001f0)
α̅(t) = cos((t / T + ϵ) / (1 + ϵ) * π / 2)^2
β = Vector{Real}(undef, T)
for t in 1:T
αₜ = α̅(t) / α̅(t - 1)
βₜ = 1 - αₜ
βₜ = min(βₘₐₓ, βₜ)
β[t] = βₜ
end
return β
end

View file

@ -0,0 +1,17 @@
"""
Linear beta schedule.
## Input
* `T::Integer`: number of timesteps
* `β₁::Real=0.0001f0`: initial (t=1) value of β
* `β₋₁::Real=0.02f0`: final (t=T) value of β
## Output
* `β::Vector{Real}`: βₜ values at each timestep t
## References
* [[2006.11239] Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239)
"""
function linear_beta_schedule(T::Integer, β₁::Real=0.0001f0, β₋₁::Real=0.02f0)
return range(start=β₁, stop=β₋₁, length=T)
end

View file

@ -0,0 +1,17 @@
"""
Scaled linear beta schedule.
## Input
* `T::Int`: number of timesteps
* `β₁::Real=0.0001f0`: initial value of β
* `β₋₁::Real=0.02f0`: final value of β
## Output
* `β::Vector{Real}`: βₜ values at each timestep t
## References
* [[2006.11239] Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239)
"""
function scaled_linear_beta_schedule(T::Integer, β₁::Real=0.0001f0, β₋₁::Real=0.02f0)
return range(start=β₁^0.5, stop=β₋₁^0.5, length=T) .^ 2
end

View file

@ -0,0 +1,21 @@
import NNlib: sigmoid
"""
Sigmoid beta schedule.
## Input
* `T::Int`: number of timesteps
* `β₁::Real=0.0001f0`: initial value of β
* `β₋₁::Real=0.02f0`: final value of β
## Output
* `β::Vector{Real}`: βₜ values at each timestep t
## References
* [[2203.02923] GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation](https://arxiv.org/abs/2203.02923)
* [github.com:MinkaiXu/GeoDiff](https://github.com/MinkaiXu/GeoDiff/blob/ea0ca48045a2f7abfccd7f0df449e45eb6eae638/models/epsnet/diffusion.py#L57)
"""
function sigmoid_beta_schedule(T::Integer, β₁::Real=0.0001f0, β₋₁::Real=0.02f0)
x = range(start=-6, stop=6, length=T)
return sigmoid(x) .* (β₋₁ - β₁) .+ β₁
end

View file

@ -0,0 +1,36 @@
"""
Rescale betas to have zero terminal Signal to Noise Ratio (SNR).
## Input
* `β::AbstractArray`: βₜ values at each timestep t
## Output
* `β::Vector{Real}`: rescaled βₜ values at each timestep t
## References
* [[2305.08891] Rescaling Diffusion Models](https://arxiv.org/abs/2305.08891) (Alg. 1)
"""
function rescale_zero_terminal_snr(β::AbstractArray)
# convert β to ⎷α̅
α = 1 .- β
α̅ = cumprod(α)
⎷α̅ = sqrt.(α̅)
# store old extrema values
⎷α̅₁ = ⎷α̅[1]
⎷α̅₋₁ = ⎷α̅[end]
# shift last timestep to zero
⎷α̅ .-= ⎷α̅₋₁
# scale so that first timestep reaches old values
⎷α̅ *= ⎷α̅₁ / (⎷α̅₁ - ⎷α̅₋₁)
# convert back ⎷α̅ to β
α̅ = ⎷α̅ .^ 2
α = α̅[2:end] ./ α̅[1:end-1]
α = vcat(α̅[1], α)
β = 1 .- α
return β
end

View file

@ -1,8 +1,9 @@
module Diffusers module Diffusers
include("BetaSchedules/BetaSchedules.jl")
# abtract types # abtract types
include("Schedulers.jl") include("Schedulers.jl")
include("BetaSchedulers.jl")
# concrete types # concrete types
include("DDPM.jl") include("DDPM.jl")