conv polish

This commit is contained in:
Mike J Innes 2017-12-15 16:24:45 +00:00
parent 386eafc443
commit 51f93d9f0e
2 changed files with 19 additions and 8 deletions

View File

@ -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

View File

@ -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)