complete SparseKernel{N}

Author: yuehhua

src/wavelet.jl

@@ -1,33 +1,44 @@
-struct SparseKernel{T,S}
-    k::Int
-    conv_blk::S
-    out_weight::T
+export
+    SparseKernel,
+    SparseKernel1D,
+    SparseKernel2D,
+    SparseKernel3D
+
+
+struct SparseKernel{N,T,S}
+    conv_blk::T
+    out_weight::S
+end
+
+function SparseKernel(filter::NTuple{N,T}, ch::Pair{S, S}; init=Flux.glorot_uniform) where {N,T,S}
+    input_dim, emb_dim = ch
+    conv = Conv(filter, input_dim=>emb_dim, relu; stride=1, pad=1, init=init)
+    W_out = Dense(emb_dim, input_dim; init=init)
+    return SparseKernel{N,typeof(conv),typeof(W_out)}(conv, W_out)
 end
 
-function SparseKernel1d(k::Int, α, c::Int=1; init=Flux.glorot_uniform)
+function SparseKernel1D(k::Int, α, c::Int=1; init=Flux.glorot_uniform)
     input_dim = c*k
     emb_dim = 128
-    conv = Conv((3,), input_dim=>emb_dim, relu; stride=1, pad=1, init=init)
-    W_out = Dense(emb_dim, input_dim; init=init)
-    return SparseKernel(k, conv, W_out)
+    return SparseKernel((3, ), input_dim=>emb_dim; init=init)
 end
 
-function SparseKernel2d(k::Int, α, c::Int=1; init=Flux.glorot_uniform)
+function SparseKernel2D(k::Int, α, c::Int=1; init=Flux.glorot_uniform)
     input_dim = c*k^2
     emb_dim = α*k^2
-    conv = Conv((3, 3), input_dim=>emb_dim, relu; stride=1, pad=1, init=init)
-    W_out = Dense(emb_dim, input_dim; init=init)
-    return SparseKernel(k, conv, W_out)
+    return SparseKernel((3, 3), input_dim=>emb_dim; init=init)
 end
 
-function SparseKernel3d(k::Int, α, c::Int=1; init=Flux.glorot_uniform)
+function SparseKernel3D(k::Int, α, c::Int=1; init=Flux.glorot_uniform)
     input_dim = c*k^2
     emb_dim = α*k^2
     conv = Conv((3, 3, 3), emb_dim=>emb_dim, relu; stride=1, pad=1, init=init)
     W_out = Dense(emb_dim, input_dim; init=init)
-    return SparseKernel(k, conv, W_out)
+    return SparseKernel{3,typeof(conv),typeof(W_out)}(conv, W_out)
 end
 
+Flux.@functor SparseKernel
+
 function (l::SparseKernel)(X::AbstractArray)
     bch_sz, _, dims_r... = reverse(size(X))
     dims = reverse(dims_r)

test/runtests.jl

@@ -1,9 +1,14 @@
 using NeuralOperators
 using Test
 using Flux
+using Zygote
+using CUDA
+
+CUDA.allowscalar(false)
 
 @testset "NeuralOperators.jl" begin
     include("fourier.jl")
+    include("wavelet.jl")
     include("model.jl")
 end
 

test/wavelet.jl

@@ -1,37 +1,53 @@
-using NeuralOperators
-using CUDA
-using Zygote
-
-CUDA.allowscalar(false)
-
-T = Float32
-k = 3
-batch_size = 32
-
-α = 4
-c = 1
-in_chs = 20
-
-
-l1 = NeuralOperators.SparseKernel1d(k, α, c)
-X = rand(T, in_chs, c*k, batch_size)
-Y = l1(X)
-gradient(x->sum(l1(x)), X)
-
-
-α = 4
-c = 3
-Nx = 5
-Ny = 7
-
-l2 = NeuralOperators.SparseKernel2d(k, α, c)
-X = rand(T, Nx, Ny, c*k^2, batch_size)
-Y = l2(X)
-gradient(x->sum(l2(x)), X)
-
-Nz = 13
-
-l3 = NeuralOperators.SparseKernel3d(k, α, c)
-X = rand(T, Nx, Ny, Nz, α*k^2, batch_size)
-Y = l3(X)
-gradient(x->sum(l3(x)), X)
+@testset "SparseKernel" begin
+    T = Float32
+    k = 3
+    batch_size = 32
+
+    @testset "1D SparseKernel" begin
+        α = 4
+        c = 1
+        in_chs = 20
+        X = rand(T, in_chs, c*k, batch_size)
+
+        l1 = SparseKernel1D(k, α, c)
+        Y = l1(X)
+        @test l1 isa SparseKernel{1}
+        @test size(Y) == size(X)
+
+        gs = gradient(()->sum(l1(X)), Flux.params(l1))
+        @test length(gs.grads) == 4
+    end
+
+    @testset "2D SparseKernel" begin
+        α = 4
+        c = 3
+        Nx = 5
+        Ny = 7
+        X = rand(T, Nx, Ny, c*k^2, batch_size)
+    
+        l2 = SparseKernel2D(k, α, c)
+        Y = l2(X)
+        @test l2 isa SparseKernel{2}
+        @test size(Y) == size(X)
+
+        gs = gradient(()->sum(l2(X)), Flux.params(l2))
+        @test length(gs.grads) == 4
+    end
+
+    @testset "3D SparseKernel" begin
+        α = 4
+        c = 3
+        Nx = 5
+        Ny = 7
+        Nz = 13
+        X = rand(T, Nx, Ny, Nz, α*k^2, batch_size)
+
+        l3 = SparseKernel3D(k, α, c)
+        Y = l3(X)
+        @test l3 isa SparseKernel{3}
+        @test size(Y) == (Nx, Ny, Nz, c*k^2, batch_size)
+
+        gs = gradient(()->sum(l3(X)), Flux.params(l3))
+        @test length(gs.grads) == 4
+    end
+end