Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Bias Fix #41

Closed
wants to merge 3 commits into from
Closed

Bias Fix #41

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
44049bf
Modify trainable function
Jun 14, 2022
7df8ad8
Remove dead code
Jun 14, 2022
d2bdaca
Add tests for addendum to FourierLayer.jl
Jun 14, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 17 additions & 7 deletions src/FourierLayer.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ struct FourierLayer{F,Tc<:Complex{<:AbstractFloat},Tr<:AbstractFloat,Bf,Bl}
bl::Bl
# Constructor for the entire fourier layer
function FourierLayer(
Wf::AbstractArray{Tc,3}, Wl::AbstractMatrix{Tr},
Wf::AbstractArray{Tc,3}, Wl::AbstractMatrix{Tr},
grid::Int, σ::F = identity,
λ::Int = 12, bf = true, bl = true) where
{F,Tc<:Complex{<:AbstractFloat},Tr<:AbstractFloat}
Expand All @@ -69,7 +69,7 @@ function FourierLayer(in::Integer, out::Integer, grid::Integer, modes = 12,
# Initialize Fourier weight matrix (only with relevant modes)
Wf = initf(in, out, modes)
# Make sure filtering works
@assert modes <= floor(Int, grid/2 + 1) "Specified modes exceed allowed maximum.
@assert modes <= floor(Int, grid/2 + 1) "Specified modes exceed allowed maximum.
The number of modes to filter must be smaller than N/2 + 1"
# Pad the fourier weight matrix with additional zeros
Wf = pad_zeros(Wf, (0, floor(Int, grid/2 + 1) - modes), dims=3)
Expand All @@ -88,10 +88,20 @@ function FourierLayer(in::Integer, out::Integer, grid::Integer, modes = 12,
end

# Only train the weight array with non-zero modes
Flux.@functor FourierLayer
Flux.trainable(a::FourierLayer) = (a.Wf[:,:,1:a.λ], a.Wl,
typeof(a.bf) != Flux.Zeros ? a.bf[:,:,1:a.λ] : nothing,
typeof(a.bl) != Flux.Zeros ? a.bl : nothing)
Flux.@functor FourierLayer

function Flux.trainable(a::FourierLayer)
trainable_params = []
push!(trainable_params,a.Wf[:,:,1:a.λ])
push!(trainable_params,a.Wl)
if a.bf!=false
push!(trainable_params,typeof(a.bf) != Flux.Zeros ? a.bf[:,:,1:a.λ] : nothing)
end
if a.bl!=false
push!(trainable_params,typeof(a.bl) != Flux.Zeros ? a.bl : nothing)
end
return trainable_params
end

# The actual layer that does stuff
function (a::FourierLayer)(x::AbstractArray)
Expand Down Expand Up @@ -133,4 +143,4 @@ function Base.show(io::IO, l::FourierLayer)
print(io, "Fourier modes: ", l.λ)
print(io, "\n")
l.σ == identity || print(io, "Activation: ", l.σ)
end
end
12 changes: 12 additions & 0 deletions test/fourierlayer.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,18 @@ using Test, Random, Flux
@test size(Flux.params(FourierLayer(128, 64, 100, 20))[4]) == (1, 64, 100)
end

# Bias tests
@testset "parameters - bias" begin
layer = FourierLayer(128,128,1024,16,gelu,bias_fourier=true,bias_linear=true)
@test length(Flux.trainable(layer)) == 4
layer = FourierLayer(128,128,1024,16,gelu,bias_fourier=true,bias_linear=false)
@test length(Flux.trainable(layer)) == 3
layer = FourierLayer(128,128,1024,16,gelu,bias_fourier=false,bias_linear=true)
@test length(Flux.trainable(layer)) == 3
layer = FourierLayer(128,128,1024,16,gelu,bias_fourier=false,bias_linear=false)
@test length(Flux.trainable(layer)) == 2
end

# Accept only Int as architecture parameters
@test_throws MethodError FourierLayer(128.5, 64, 100, 20)
@test_throws MethodError FourierLayer(128.5, 64, 100, 20, tanh)
Expand Down