Merge pull request #14 from yuehhua/spectralconv

Refactor SpectralConv

Author: foldfelis

README.md

@@ -24,7 +24,7 @@ Currently, `FourierOperator` is provided in this work.
 
 ## Usage
 
-```
+```julia
 function FourierNeuralOperator()
     modes = (16, )
     ch = 64 => 64
@@ -48,13 +48,13 @@ end
 
 Or you can just call:
 
-```
+```julia
 fno = FourierNeuralOperator()
 ```
 
 And then train as a Flux model.
 
-```
+```julia
 loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- fno(𝐱)) / size(𝐱)[end]
 opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
 Flux.@epochs 50 Flux.train!(loss, params(m), data, opt)

src/fourier.jl

@@ -1,19 +1,8 @@
 export
     SpectralConv,
-    SpectralConvPerm,
     FourierOperator
 
-abstract type AbstractSpectralConv{N, T, S, F} end
-
-struct SpectralConv{N, T, S, F} <: AbstractSpectralConv{N, T, S, F}
-    weight::T
-    in_channel::S
-    out_channel::S
-    modes::NTuple{N, S}
-    σ::F
-end
-
-struct SpectralConvPerm{N, T, S, F} <: AbstractSpectralConv{N, T, S, F}
+struct SpectralConv{P, N, T, S, F}
     weight::T
     in_channel::S
     out_channel::S
@@ -33,21 +22,21 @@ end
 * `modes`: The Fourier modes to be preserved.
 * `σ`: Activation function.
 * `permuted`: Whether the dim is permuted. If `permuted=true`, layer accepts
-    data in the order of `(..., ch, batch)`, otherwise the order is `(ch, ..., batch)`.
+    data in the order of `(ch, ..., batch)`, otherwise the order is `(..., ch, batch)`.
 
 ## Example
 
 ```jldoctest
 julia> SpectralConv(2=>5, (16, ))
-SpectralConv(2 => 5, (16,), σ=identity)
+SpectralConv(2 => 5, (16,), σ=identity, permuted=false)
 
 julia> using Flux
 
 julia> SpectralConv(2=>5, (16, ), relu)
-SpectralConv(2 => 5, (16,), σ=relu)
+SpectralConv(2 => 5, (16,), σ=relu, permuted=false)
 
 julia> SpectralConv(2=>5, (16, ), relu, permuted=true)
-SpectralConvPerm(2 => 5, (16,), σ=relu)
+SpectralConv(2 => 5, (16,), σ=relu, permuted=true)
 ```
 """
 function SpectralConv(
@@ -61,23 +50,23 @@ function SpectralConv(
     in_chs, out_chs = ch
     scale = one(T) / (in_chs * out_chs)
     weights = scale * init(out_chs, in_chs, prod(modes))
+    W = typeof(weights)
+    F = typeof(σ)
 
-    L = permuted ? SpectralConvPerm : SpectralConv
-
-    return L(weights, in_chs, out_chs, modes, σ)
+    return SpectralConv{permuted,N,W,S,F}(weights, in_chs, out_chs, modes, σ)
 end
 
 Flux.@functor SpectralConv
-Flux.@functor SpectralConvPerm
 
-Base.ndims(::AbstractSpectralConv{N}) where {N} = N
+Base.ndims(::SpectralConv{P,N}) where {P,N} = N
+
+permuted(::SpectralConv{P}) where {P} = P
 
-function Base.show(io::IO, l::AbstractSpectralConv)
-    T = (l isa SpectralConv) ? SpectralConv : SpectralConvPerm
-    print(io, "$(string(T))($(l.in_channel) => $(l.out_channel), $(l.modes), σ=$(string(l.σ)))")
+function Base.show(io::IO, l::SpectralConv{P}) where {P}
+    print(io, "SpectralConv($(l.in_channel) => $(l.out_channel), $(l.modes), σ=$(string(l.σ)), permuted=$P)")
 end
 
-function spectral_conv(m::AbstractSpectralConv, 𝐱::AbstractArray)
+function spectral_conv(m::SpectralConv, 𝐱::AbstractArray)
     n_dims = ndims(𝐱)
 
     𝐱_fft = fft(Zygote.hook(real, 𝐱), 1:ndims(m)) # [x, in_chs, batch]
@@ -90,22 +79,28 @@ function spectral_conv(m::AbstractSpectralConv, 𝐱::AbstractArray)
     return m.σ.(𝐱_ifft)
 end
 
-function (m::SpectralConv)(𝐱)
+function (m::SpectralConv{false})(𝐱)
     𝐱ᵀ = permutedims(𝐱, (ntuple(i->i+1, ndims(m))..., 1, ndims(m)+2)) # [x, in_chs, batch] <- [in_chs, x, batch]
     𝐱_out = spectral_conv(m, 𝐱ᵀ) # [x, out_chs, batch]
     𝐱_outᵀ = permutedims(𝐱_out, (ndims(m)+1, 1:ndims(m)..., ndims(m)+2)) # [out_chs, x, batch] <- [x, out_chs, batch]
 
     return 𝐱_outᵀ
 end
 
-function (m::SpectralConvPerm)(𝐱)
+function (m::SpectralConv{true})(𝐱)
     return spectral_conv(m, 𝐱) # [x, out_chs, batch]
 end
 
 ############
 # operator #
 ############
 
+struct FourierOperator{L, C, F}
+    linear::L
+    conv::C
+    σ::F
+end
+
 """
     FourierOperator(ch, modes, σ=identity; permuted=false)
 
@@ -115,42 +110,21 @@ end
 * `modes`: The Fourier modes to be preserved for spectral convolution.
 * `σ`: Activation function.
 * `permuted`: Whether the dim is permuted. If `permuted=true`, layer accepts
-    data in the order of `(..., ch, batch)`, otherwise the order is `(ch, ..., batch)`.
+    data in the order of `(ch, ..., batch)`, otherwise the order is `(..., ch, batch)`.
 
 ## Example
 
 ```jldoctest
 julia> FourierOperator(2=>5, (16, ))
-Chain(
-  Parallel(
-    +,
-    Dense(2, 5),                        # 15 parameters
-    SpectralConv(2 => 5, (16,), σ=identity),  # 160 parameters
-  ),
-  NeuralOperators.var"#activation_func#14"{typeof(identity)}(identity),
-)                   # Total: 3 arrays, 175 parameters, 1.668 KiB.
+FourierOperator(2 => 5, (16,), σ=identity, permuted=false)
 
 julia> using Flux
 
 julia> FourierOperator(2=>5, (16, ), relu)
-Chain(
-  Parallel(
-    +,
-    Dense(2, 5),                        # 15 parameters
-    SpectralConv(2 => 5, (16,), σ=identity),  # 160 parameters
-  ),
-  NeuralOperators.var"#activation_func#14"{typeof(relu)}(NNlib.relu),
-)                   # Total: 3 arrays, 175 parameters, 1.668 KiB.
+FourierOperator(2 => 5, (16,), σ=relu, permuted=false)
 
 julia> FourierOperator(2=>5, (16, ), relu, permuted=true)
-Chain(
-  Parallel(
-    +,
-    Conv((1,), 2 => 5),                 # 15 parameters
-    SpectralConvPerm(2 => 5, (16,), σ=identity),  # 160 parameters
-  ),
-  NeuralOperators.var"#activation_func#14"{typeof(relu)}(NNlib.relu),
-)                   # Total: 3 arrays, 175 parameters, 1.871 KiB.
+FourierOperator(2 => 5, (16,), σ=relu, permuted=true)
 ```
 """
 function FourierOperator(
@@ -159,15 +133,23 @@ function FourierOperator(
     σ=identity;
     permuted=false
 ) where {S<:Integer, N}
-    short_cut = permuted ? Conv(Tuple(ones(Int, length(modes))), ch) : Dense(ch.first, ch.second)
-    activation_func(x) = σ.(x)
+    linear = permuted ? Conv(Tuple(ones(Int, length(modes))), ch) : Dense(ch.first, ch.second)
+    conv = SpectralConv(ch, modes; permuted=permuted)
 
-    return Chain(
-        Parallel(+, short_cut, SpectralConv(ch, modes, permuted=permuted)),
-        activation_func
-    )
+    return FourierOperator(linear, conv, σ)
+end
+
+Flux.@functor FourierOperator
+
+function Base.show(io::IO, l::FourierOperator)
+    print(io, "FourierOperator($(l.conv.in_channel) => $(l.conv.out_channel), $(l.conv.modes), σ=$(string(l.σ)), permuted=$(permuted(l.conv)))")
 end
 
+function (m::FourierOperator)(𝐱)
+    return m.σ.(m.linear(𝐱) + m.conv(𝐱))
+end
+
+
 #########
 # utils #
 #########

test/fourier.jl

@@ -1,4 +1,4 @@
-@testset "SpectralConv1d" begin
+@testset "1D SpectralConv" begin
     modes = (16, )
     ch = 64 => 128
 
@@ -7,7 +7,7 @@
         SpectralConv(ch, modes)
     )
     @test ndims(SpectralConv(ch, modes)) == 1
-    @test repr(SpectralConv(ch, modes)) == "SpectralConv(64 => 128, (16,), σ=identity)"
+    @test repr(SpectralConv(ch, modes)) == "SpectralConv(64 => 128, (16,), σ=identity, permuted=false)"
 
     𝐱, _ = get_burgers_data(n=5)
     @test size(m(𝐱)) == (128, 1024, 5)
@@ -17,7 +17,7 @@
     Flux.train!(loss, params(m), data, Flux.ADAM())
 end
 
-@testset "SpectralConvPerm1d" begin
+@testset "permuted 1D SpectralConv" begin
     modes = (16, )
     ch = 64 => 128
 
@@ -26,6 +26,7 @@ end
         SpectralConv(ch, modes, permuted=true)
     )
     @test ndims(SpectralConv(ch, modes, permuted=true)) == 1
+    @test repr(SpectralConv(ch, modes, permuted=true)) == "SpectralConv(64 => 128, (16,), σ=identity, permuted=true)"
 
     𝐱, _ = get_burgers_data(n=5)
     𝐱 = permutedims(𝐱, (2, 1, 3))
@@ -36,14 +37,15 @@ end
     Flux.train!(loss, params(m), data, Flux.ADAM())
 end
 
-@testset "FourierOperator1d" begin
+@testset "1D FourierOperator" begin
     modes = (16, )
     ch = 64 => 128
 
     m = Chain(
         Dense(2, 64),
         FourierOperator(ch, modes)
     )
+    @test repr(FourierOperator(ch, modes)) == "FourierOperator(64 => 128, (16,), σ=identity, permuted=false)"
 
     𝐱, _ = get_burgers_data(n=5)
     @test size(m(𝐱)) == (128, 1024, 5)
@@ -53,14 +55,15 @@ end
     Flux.train!(loss, params(m), data, Flux.ADAM())
 end
 
-@testset "FourierOperatorPerm1d" begin
+@testset "permuted 1D FourierOperator" begin
     modes = (16, )
     ch = 64 => 128
 
     m = Chain(
         Conv((1, ), 2=>64),
         FourierOperator(ch, modes, permuted=true)
     )
+    @test repr(FourierOperator(ch, modes, permuted=true)) == "FourierOperator(64 => 128, (16,), σ=identity, permuted=true)"
 
     𝐱, _ = get_burgers_data(n=5)
     𝐱 = permutedims(𝐱, (2, 1, 3))
@@ -71,7 +74,7 @@ end
     Flux.train!(loss, params(m), data, Flux.ADAM())
 end
 
-@testset "SpectralConv2d" begin
+@testset "2D SpectralConv" begin
     modes = (16, 16)
     ch = 64 => 64
 
@@ -89,7 +92,7 @@ end
     Flux.train!(loss, params(m), data, Flux.ADAM())
 end
 
-@testset "SpectralConvPerm2d" begin
+@testset "permuted 2D SpectralConv" begin
     modes = (16, 16)
     ch = 64 => 64
 
@@ -108,7 +111,7 @@ end
     Flux.train!(loss, params(m), data, Flux.ADAM())
 end
 
-@testset "FourierOperator2d" begin
+@testset "2D FourierOperator" begin
     modes = (16, 16)
     ch = 64 => 64
 
@@ -125,7 +128,7 @@ end
     Flux.train!(loss, params(m), data, Flux.ADAM())
 end
 
-@testset "FourierOperatorPerm2d" begin
+@testset "permuted 2D FourierOperator" begin
     modes = (16, 16)
     ch = 64 => 64