Merge pull request #13 from yuehhua/develop

Fix multihead issue
yuehhua · Jul 21, 2022 · 29855d1 · 29855d1
2 parents a970d4b + 1d1908b
commit 29855d1
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 22 deletions.
diff --git a/src/functional.jl b/src/functional.jl
@@ -11,7 +11,7 @@ Hopfield layer forward function.
 - `key::AbstractArray`:
 - `value::AbstractArray`:
 - `max_iter`: `-1` for iteration until stable.
-- `ϵ`: 
+- `ϵ`: The error tolerance for convergence.
 """
 function hopfield_forward(Qt, Kt, Vt, out_proj, dropout, heads::Int,
     β::AbstractArray, query::AbstractArray, key::AbstractArray, value::AbstractArray,
@@ -25,11 +25,10 @@ function hopfield_forward(Qt, Kt, Vt, out_proj, dropout, heads::Int,
 
     Â = attention_prob(Q, K, β)
     Â = multiple_updates(Â, Q, K, β, heads, max_iter, ϵ)
-    Â = move_heads_to_first(Â, heads)
-    V = move_heads_to_first(V, heads)
 
     V = dropout(V)
-    attn_out = batched_mul(V, batched_transpose(Â))
+    attn_out = batched_innerprod(V, Â, dims=2)
+    attn_out = move_heads_to_first(attn_out, heads)
     return out_proj(attn_out)
 end
 

diff --git a/test/cuda.jl b/test/cuda.jl
@@ -18,7 +18,7 @@
         l = HopfieldCore(emb_dim, heads; kdim=kdim, vdim=vdim,
             head_dim=head_dim, pattern_dim=pattern_dim) |> gpu
         Y = l(Q, K, V)
-        @test size(Y) == (emb_dim, heads*target_len, batch_size)
+        @test size(Y) == (emb_dim, target_len, batch_size)
 
         g = Zygote.gradient(() -> sum(l(Q, K, V)), Flux.params(l))
         @test length(g.grads) == 6
@@ -28,7 +28,7 @@
         l = HopfieldCore(emb_dim, heads; kdim=kdim, vdim=vdim,
             head_dim=head_dim, pattern_dim=pattern_dim, enable_out_proj=false) |> gpu
         Y = l(Q, K, V)
-        @test size(Y) == (heads*pattern_dim, heads*target_len, batch_size)
+        @test size(Y) == (heads*pattern_dim, target_len, batch_size)
 
         g = Zygote.gradient(() -> sum(l(Q, K, V)), Flux.params(l))
         @test length(g.grads) == 4
@@ -38,17 +38,17 @@
         l = HopfieldCore(emb_dim, heads; kdim=kdim, vdim=kdim,
             head_dim=head_dim, pattern_dim=pattern_dim) |> gpu
         Y = l(Q, nothing, nothing)
-        @test size(Y) == (emb_dim, heads*target_len, batch_size)
+        @test size(Y) == (emb_dim, target_len, batch_size)
 
         g = Zygote.gradient(() -> sum(l(Q, nothing, nothing)), Flux.params(l))
-        @test length(g.grads) == 8
+        @test length(g.grads) == 10
     end
 
     @testset "static Q" begin
         l = HopfieldCore(emb_dim, heads; kdim=kdim, vdim=vdim,
             head_dim=head_dim, pattern_dim=pattern_dim) |> gpu
         Y = l(nothing, K, V)
-        @test size(Y) == (emb_dim, heads*emb_dim, batch_size)
+        @test size(Y) == (emb_dim, emb_dim, batch_size)
 
         g = Zygote.gradient(() -> sum(l(nothing, K, V)), Flux.params(l))
         @test length(g.grads) == 8

diff --git a/test/layer.jl b/test/layer.jl
@@ -31,7 +31,7 @@
         @test size(l.linear_v(V)) == (heads*pattern_dim, source_len, batch_size)
 
         Y = l(Q, K, V)
-        @test size(Y) == (emb_dim, heads*target_len, batch_size)
+        @test size(Y) == (emb_dim, target_len, batch_size)
 
         g = Zygote.gradient(() -> sum(l(Q, K, V)), Flux.params(l))
         @test length(g.grads) == 6
@@ -41,7 +41,7 @@
         l = HopfieldCore(emb_dim, heads; kdim=kdim, vdim=vdim,
             head_dim=head_dim, pattern_dim=pattern_dim, enable_out_proj=false)
         Y = l(Q, K, V)
-        @test size(Y) == (heads*pattern_dim, heads*target_len, batch_size)
+        @test size(Y) == (heads*pattern_dim, target_len, batch_size)
 
         g = Zygote.gradient(() -> sum(l(Q, K, V)), Flux.params(l))
         @test length(g.grads) == 4
@@ -51,17 +51,17 @@
         l = HopfieldCore(emb_dim, heads; kdim=kdim, vdim=kdim,
             head_dim=head_dim, pattern_dim=pattern_dim)
         Y = l(Q, nothing, nothing)
-        @test size(Y) == (emb_dim, heads*target_len, batch_size)
+        @test size(Y) == (emb_dim, target_len, batch_size)
 
         g = Zygote.gradient(() -> sum(l(Q, nothing, nothing)), Flux.params(l))
-        @test length(g.grads) == 8
+        @test length(g.grads) == 10
     end
 
     @testset "static Q" begin
         l = HopfieldCore(emb_dim, heads; kdim=kdim, vdim=vdim,
             head_dim=head_dim, pattern_dim=pattern_dim)
         Y = l(nothing, K, V)
-        @test size(Y) == (emb_dim, heads*emb_dim, batch_size)
+        @test size(Y) == (emb_dim, emb_dim, batch_size)
 
         g = Zygote.gradient(() -> sum(l(nothing, K, V)), Flux.params(l))
         @test length(g.grads) == 8
@@ -81,7 +81,7 @@
         @test Hopfields.pattern_projection_dim(l) == pattern_projection_dim
 
         Y = l(Q, K, V)
-        @test size(Y) == (out_channel, heads*target_len, batch_size)
+        @test size(Y) == (out_channel, target_len, batch_size)
 
         g = Zygote.gradient(() -> sum(l(Q, K, V)), Flux.params(l))
         @test length(g.grads) == 6
@@ -100,10 +100,10 @@
         @test Hopfields.pattern_projection_dim(l) == stored_pattern_dim
 
         Y = l(Q, nothing, nothing)
-        @test size(Y) == (out_channel, heads*target_len, batch_size)
+        @test size(Y) == (out_channel, target_len, batch_size)
 
         g = Zygote.gradient(() -> sum(l(Q, nothing, nothing)), Flux.params(l))
-        @test length(g.grads) == 8
+        @test length(g.grads) == 10
     end
 
     @testset "HopfieldPooling" begin
@@ -120,7 +120,7 @@
         @test Hopfields.pattern_projection_dim(l) == pattern_projection_dim
 
         Y = l(nothing, K, V)
-        @test size(Y) == (out_channel, heads*emb_dim, batch_size)
+        @test size(Y) == (out_channel, emb_dim, batch_size)
 
         g = Zygote.gradient(() -> sum(l(nothing, K, V)), Flux.params(l))
         @test length(g.grads) == 8
@@ -132,7 +132,7 @@
             pattern_projection_dim=pattern_projection_dim, max_iter=3)
 
         Y = l(Q, K, V)
-        @test size(Y) == (out_channel, heads*target_len, batch_size)
+        @test size(Y) == (out_channel, target_len, batch_size)
 
         g = Zygote.gradient(() -> sum(l(Q, K, V)), Flux.params(l))
         @test length(g.grads) == 6
@@ -143,10 +143,10 @@
             stored_pattern_dim=stored_pattern_dim, max_iter=3)
 
         Y = l(Q, nothing, nothing)
-        @test size(Y) == (out_channel, heads*target_len, batch_size)
+        @test size(Y) == (out_channel, target_len, batch_size)
 
         g = Zygote.gradient(() -> sum(l(Q, nothing, nothing)), Flux.params(l))
-        @test length(g.grads) == 20
+        @test length(g.grads) == 22
     end
 
     @testset "HopfieldPooling with multiple updates" begin
@@ -155,7 +155,7 @@
             pattern_projection_dim=pattern_projection_dim, max_iter=3)
 
         Y = l(nothing, K, V)
-        @test size(Y) == (out_channel, heads*emb_dim, batch_size)
+        @test size(Y) == (out_channel, emb_dim, batch_size)
 
         g = Zygote.gradient(() -> sum(l(nothing, K, V)), Flux.params(l))
         @test length(g.grads) == 8