Fix reshape
This commit is contained in:
parent
e2ae8b4e8d
commit
d6aacf4135
@ -121,13 +121,15 @@ function cudnnBNForward!(y::CuArray{T}, g::CuArray{T}, b::CuArray{T}, x::CuArray
|
||||
end
|
||||
end
|
||||
|
||||
∇batchnorm(g::CuArray{T}, b::CuArray{T}, x::CuArray{T, 2}, dy::CuArray{T},
|
||||
function ∇batchnorm(g::CuArray{T}, b::CuArray{T}, x::CuArray{T, 2}, dy::CuArray{T, 2},
|
||||
running_mean::CuArray{T}, running_var::CuArray{T}, momentum;
|
||||
cache = nothing, eps = T(1e-5), alpha = T(1),
|
||||
beta = T(0), training = true) where T<:Union{Float32, Float64} =
|
||||
∇batchnorm(g, b, reshape(x, 1, 1, size(x, 1), size(x, 2)), dy,
|
||||
running_mean, running_var, momentum, cache = cache, eps = eps, alpha = alpha, beta = beta,
|
||||
training = training)
|
||||
beta = T(0), training = true) where T<:Union{Float32, Float64}
|
||||
dg, db, dx = ∇batchnorm(g, b, reshape(x, 1, 1, size(x, 1), size(x, 2)), reshape(dy, 1, 1, size(dy, 1),
|
||||
size(dy, 2)), running_mean, running_var, momentum, cache = cache, eps = eps,
|
||||
alpha = alpha, beta = beta, training = training)
|
||||
(dg, db, dropdims(dx, dims = (1, 2)))
|
||||
end
|
||||
|
||||
function ∇batchnorm(g::CuArray{T}, b::CuArray{T}, x::CuArray{T}, dy::CuArray{T},
|
||||
running_mean::CuArray{T}, running_var::CuArray{T}, momentum;
|
||||
|
Loading…
Reference in New Issue
Block a user