conv polish
This commit is contained in:
parent
386eafc443
commit
51f93d9f0e
@ -10,5 +10,11 @@ Conv2D(k::NTuple{2,Integer}, ch::Pair{<:Integer,<:Integer}, σ = identity;
|
||||
|
||||
Flux.treelike(Conv2D)
|
||||
|
||||
# (c::Conv2D)(x) = c.σ.(conv2d(x, c.weight, stride = c.stride))
|
||||
(c::Conv2D)(x) = c.σ.(conv2d(x, c.weight))
|
||||
(c::Conv2D)(x) = c.σ.(conv2d(x, c.weight, stride = c.stride))
|
||||
|
||||
function Base.show(io::IO, l::Conv2D)
|
||||
print(io, "Conv2D((", size(l.weight, 1), ", ", size(l.weight, 2), ")")
|
||||
print(io, ", ", size(l.weight, 3), "=>", size(l.weight, 4))
|
||||
l.σ == identity || print(io, ", ", l.σ)
|
||||
print(io, ")")
|
||||
end
|
||||
|
@ -136,13 +136,18 @@ softmax(xs::TrackedArray) = TrackedArray(Call(softmax, xs))
|
||||
|
||||
back(::typeof(softmax), Δ, xs) = @back(xs, ∇softmax(Δ, data(xs)))
|
||||
|
||||
conv2d(x::TrackedArray{<:Any,4}, w::TrackedArray{<:Any,4}) = TrackedArray(Call(conv2d, x, w))
|
||||
conv2d(x::AbstractArray{<:Any,4}, w::TrackedArray{<:Any,4}) = TrackedArray(Call(conv2d, x, w))
|
||||
conv2d(x::TrackedArray{<:Any,4}, w::AbstractArray{<:Any,4}) = TrackedArray(Call(conv2d, x, w))
|
||||
_conv2d(x, w, stride) = conv2d(x, w, stride = stride)
|
||||
|
||||
function back(::typeof(conv2d), Δ, x, w)
|
||||
@back(x, NNlib.conv2d_grad_x(data(x), data(w), Δ))
|
||||
@back(w, NNlib.conv2d_grad_w(data(x), data(w), Δ))
|
||||
conv2d(x::TrackedArray{<:Any,4}, w::TrackedArray{<:Any,4}; stride = 1) =
|
||||
TrackedArray(Call(_conv2d, x, w, stride))
|
||||
conv2d(x::AbstractArray{<:Any,4}, w::TrackedArray{<:Any,4}; stride = 1) =
|
||||
TrackedArray(Call(_conv2d, x, w, stride))
|
||||
conv2d(x::TrackedArray{<:Any,4}, w::AbstractArray{<:Any,4}; stride = 1) =
|
||||
TrackedArray(Call(_conv2d, x, w, stride))
|
||||
|
||||
function back(::typeof(_conv2d), Δ, x, w, stride)
|
||||
@back(x, NNlib.conv2d_grad_x(data(x), data(w), Δ; stride = stride))
|
||||
@back(w, NNlib.conv2d_grad_w(data(x), data(w), Δ; stride = stride))
|
||||
end
|
||||
|
||||
_pool(x, k, mode) = pool(x, window = k, mode = mode)
|
||||
|
Loading…
Reference in New Issue
Block a user