mirror of
https://github.com/Laurent2916/Diffusers.jl.git
synced 2024-09-18 18:45:28 +00:00
🚚 (BetaSchedulers) move beta variance schedules in their own module
This commit is contained in:
parent
f6ecfdbfda
commit
4dfa3c92dc
|
@ -4,7 +4,6 @@
|
|||
|
||||
[🤗 diffusers](https://github.com/huggingface/diffusers/), but in Julia (and with worse code) !
|
||||
|
||||
|
||||
## Credits
|
||||
|
||||
- [@huggingface](https://github.com/huggingface) for their amazing [🤗 diffusers](https://github.com/huggingface/diffusers/) library.
|
||||
|
|
|
@ -1,4 +1,5 @@
|
|||
import Diffusers
|
||||
import Diffusers.BetaSchedules: cosine_beta_schedule, rescale_zero_terminal_snr
|
||||
using Flux
|
||||
using Random
|
||||
using Plots
|
||||
|
@ -42,8 +43,8 @@ scatter(data[1, :], data[2, :],
|
|||
num_timesteps = 100
|
||||
scheduler = Diffusers.DDPM(
|
||||
Vector{Float64},
|
||||
Diffusers.rescale_zero_terminal_snr(
|
||||
Diffusers.cosine_beta_schedule(num_timesteps, 0.999f0, 0.001f0)
|
||||
rescale_zero_terminal_snr(
|
||||
cosine_beta_schedule(num_timesteps, 0.999f0, 0.001f0)
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
@ -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
|
21
src/BetaSchedules/BetaSchedules.jl
Normal file
21
src/BetaSchedules/BetaSchedules.jl
Normal 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
|
30
src/BetaSchedules/Cosine.jl
Normal file
30
src/BetaSchedules/Cosine.jl
Normal 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
|
17
src/BetaSchedules/Linear.jl
Normal file
17
src/BetaSchedules/Linear.jl
Normal 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
|
17
src/BetaSchedules/ScaledLinear.jl
Normal file
17
src/BetaSchedules/ScaledLinear.jl
Normal 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
|
21
src/BetaSchedules/Sigmoid.jl
Normal file
21
src/BetaSchedules/Sigmoid.jl
Normal 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
|
36
src/BetaSchedules/ZeroSNR.jl
Normal file
36
src/BetaSchedules/ZeroSNR.jl
Normal 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
|
|
@ -1,8 +1,9 @@
|
|||
module Diffusers
|
||||
|
||||
include("BetaSchedules/BetaSchedules.jl")
|
||||
|
||||
# abtract types
|
||||
include("Schedulers.jl")
|
||||
include("BetaSchedulers.jl")
|
||||
|
||||
# concrete types
|
||||
include("DDPM.jl")
|
||||
|
|
Loading…
Reference in a new issue