From 269d8f36b9a766301197d0176a405acdd841b890 Mon Sep 17 00:00:00 2001 From: Mike J Innes Date: Mon, 18 Dec 2017 18:05:38 +0000 Subject: [PATCH] conv padding --- src/layers/conv.jl | 7 ++++--- src/tracker/lib.jl | 21 +++++++++++---------- 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/src/layers/conv.jl b/src/layers/conv.jl index 82d90029..d73d1ad9 100644 --- a/src/layers/conv.jl +++ b/src/layers/conv.jl @@ -2,15 +2,16 @@ struct Conv2D{F,A} σ::F weight::A stride::Int + pad::Int end Conv2D(k::NTuple{2,Integer}, ch::Pair{<:Integer,<:Integer}, σ = identity; - init = initn, stride = 1) = - Conv2D(σ, param(initn(k..., ch...)), stride) + init = initn, stride = 1, pad = 0) = + Conv2D(σ, param(initn(k..., ch...)), stride, pad) Flux.treelike(Conv2D) -(c::Conv2D)(x) = c.σ.(conv2d(x, c.weight, stride = c.stride)) +(c::Conv2D)(x) = c.σ.(conv2d(x, c.weight, stride = c.stride, padding = c.pad)) function Base.show(io::IO, l::Conv2D) print(io, "Conv2D((", size(l.weight, 1), ", ", size(l.weight, 2), ")") diff --git a/src/tracker/lib.jl b/src/tracker/lib.jl index 580992ef..2dc25e52 100644 --- a/src/tracker/lib.jl +++ b/src/tracker/lib.jl @@ -136,18 +136,19 @@ softmax(xs::TrackedArray) = TrackedArray(Call(softmax, xs)) back(::typeof(softmax), Δ, xs) = @back(xs, ∇softmax(Δ, data(xs))) -_conv2d(x, w, stride) = conv2d(x, w, stride = stride) +# TODO: can store kwargs efficiently in namedtuples +_conv2d(x, w, stride, pad) = conv2d(x, w, stride = stride, padding = pad) -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)) +conv2d(x::TrackedArray{<:Any,4}, w::TrackedArray{<:Any,4}; stride = 1, padding = 0) = + TrackedArray(Call(_conv2d, x, w, stride, padding)) +conv2d(x::AbstractArray{<:Any,4}, w::TrackedArray{<:Any,4}; stride = 1, padding = 0) = + TrackedArray(Call(_conv2d, x, w, stride, padding)) +conv2d(x::TrackedArray{<:Any,4}, w::AbstractArray{<:Any,4}; stride = 1, padding = 0) = + TrackedArray(Call(_conv2d, x, w, stride, padding)) -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)) +function back(::typeof(_conv2d), Δ, x, w, stride, pad) + @back(x, NNlib.conv2d_grad_x(data(x), data(w), Δ; stride = stride, padding = pad)) + @back(w, NNlib.conv2d_grad_w(data(x), data(w), Δ; stride = stride, padding = pad)) end _pool(x, k, mode) = pool(x, window = k, mode = mode)