Skip to content

Added nnlib layers and test cases, added upsampling layers #38

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bors merged 2 commits into from
Sep 8, 2020
Merged

Added nnlib layers and test cases, added upsampling layers #38

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bf63ed1
Added nnlib layers and test cases, added upsampling layers
BjornNorell Aug 24, 2020
1a45e41
Fixed broken upsampling layers
BjornNorell Aug 26, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 38 additions & 13 deletions src/nnlib.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,18 +2,31 @@ using NNlib
using NNlib: expand
using NNlib: PoolDims

import NNlib: conv
import NNlib: conv, depthwiseconv

function NNlib.conv(x::Tensor{xT, N}, w::Tensor, b::Tensor{T}, cdims::DenseConvDims{M,K,C_in,C_out,S,P,D,F};
stride = 1, pad = 0, dilation = 1) where {T,N, xT, M,K,C_in,C_out,S,P,D,F}

op = conv2d(x, w, b, stride = collect(S), padding = [P[1];P[3]], dilation = collect(dilation))
function NNlib.conv(x::Tensor{xT, N}, w::Tensor, b::Tensor{T},
cdims::DenseConvDims{M,K,C_in,C_out,S,P,D,F}) where {T,N,xT,M,K,C_in,C_out,S,P,D,F}
op = conv2d(x, w, b, stride = collect(S), padding = [P[1];P[3]], dilation = collect(D))
op
end

function NNlib.conv(x::Tensor, w::Tensor, cdims::DenseConvDims; stride = 1, pad = 0, dilation = 1)
function NNlib.conv(x::Tensor, w::Tensor, cdims::DenseConvDims)
b = zeros(Tensor{Float32}, size(w)[end], dev = on(w))
op = conv(x, w, b, cdims, stride = stride, pad = pad, dilation = dilation)
op = conv(x, w, b, cdims)
op
end

function NNlib.depthwiseconv(x::Tensor{xT, N}, w::Tensor, b::Tensor{T};
stride = 1, pad = 0, dilation = 1) where {T, N, xT}
op = _depthwise_conv2d(x, w, b, stride = collect(stride), padding = collect(pad),
dilation = collect(dilation))
op
end

function NNlib.depthwiseconv(x::Tensor, w::Tensor; stride = 1, pad = 0, dilation = 1)
b = zeros(Tensor{Float32}, size(w)[end], dev = on(w))
op = depthwiseconv(x, w, b, stride = collect(stride), pad = collect(pad),
dilation = collect(dilation))
op
end

Expand All @@ -38,8 +51,15 @@ function NNlib.sigmoid(t::Tensor{T,N}) where {T,N}
Tensor{T,N}(ptr[], on(t))
end

function NNlib.tanh(t::Tensor{T,N}) where {T,N}
ptr = Ref(Ptr{Cvoid}())

atg_tanh(ptr, t.ptr)
Tensor{T,N}(ptr[], on(t))
end

function NNlib.softmax(t::Tensor{T,N}; dims = 1) where {T,N}
_softmax(t, N - dims, options[T])
_softmax(t, dims, options[T])
end

function NNlib.∇softmax(Δ, xs::Tensor; dims = 1)
Expand All @@ -48,15 +68,20 @@ function NNlib.∇softmax(Δ, xs::Tensor; dims = 1)
sf .* (t .- sum(t .* sf, dims = dims))
end

function NNlib.meanpool(t::Tensor, pdims::PoolDims{N,K,S,P,D}; kw...) where {N,K,S,P,D}
function NNlib.meanpool(t::Tensor, pdims::PoolDims{N,K,S,P,D}) where {N,K,S,P,D}
ks = collect(NNlib.kernel_size(pdims))
stride = collect(S)
pad = [P[1];P[3]]
op_sz = NNlib.output_size(pdims)
padding = [P[1];P[3]]
# op_sz = NNlib.output_size(pdims)

_meanpool(t, ks, stride, pad, op_sz)
_meanpool(t, ks, stride=stride, padding=padding)
end

function NNlib.maxpool(t::Tensor, pdims::PoolDims{N,K,S,P,D}) where {N,K,S,P,D}
_maxpool(t, pdims)
ks = collect(NNlib.kernel_size(pdims))
stride = collect(S)
padding = [P[1];P[3]]
dilation = collect(D)

_maxpool(t, ks, stride=stride, padding=padding, dilation=dilation)
end
164 changes: 139 additions & 25 deletions src/ops.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,39 +82,102 @@ end

# TODO: Use a macro to generate wrappers
function conv2d(input::Tensor{T}, filter::Tensor{T,N}, bias::Tensor{T};
stride = [1],
padding = [0],
dilation = [1],
groups = 1) where {T,N}
stride = [1],
padding = [0],
dilation = [1],
groups = 1) where {T,N}

ptr = Ref(Ptr{Cvoid}())

atg_conv2d(ptr, input.ptr, filter.ptr, bias.ptr,
stride, length(stride),
padding, length(padding),
dilation, length(dilation),
groups)
reverse(stride), length(stride),
reverse(padding), length(padding),
reverse(dilation), length(dilation),
groups)

Tensor{T,N}(ptr[], on(input))
end

function conv_transpose_2d(input::Tensor{T}, filter::Tensor{T,N}, bias::Tensor{T};
stride = [1],
padding = [0],
output_padding = [0],
dilation = [1],
groups = 1) where {T,N}

ptr = Ref(Ptr{Cvoid}())

atg_conv_transpose2d(ptr, input.ptr, filter.ptr, bias.ptr,
reverse(stride), length(stride),
reverse(padding), length(padding),
reverse(output_padding), length(output_padding),
groups,
reverse(dilation), length(dilation))

Tensor{T,N}(ptr[], on(input))
end

function _depthwise_conv2d(input::Tensor{T}, filter::Tensor{T,N}, bias::Tensor{T};
stride = [1],
padding = [0],
dilation = [1]) where {T,N}

# When groups == in_channels and out_channels == K * in_channels, where K is a positive integer,
# this operation is also termed in literature as depthwise convolution.

c_in = size(input)[end - 1] # number of input channels
c_out = size(filter)[end] # number of output channels
@assert mod(c_in, c_out) == 0 "Invalid kernel size for depthwise convolution"

groups = c_in
ptr = Ref(Ptr{Cvoid}())

atg_conv2d(ptr, input.ptr, filter.ptr, bias.ptr,
reverse(stride), length(stride),
reverse(padding), length(padding),
reverse(dilation), length(dilation),
groups)

Tensor{T,N}(ptr[], on(input))
end

function _softmax(input::Tensor{T,N}, dims = 1, dtype = options[T]) where {T,N}
ptr = Ref(Ptr{Cvoid}())

atg_softmax(ptr, input.ptr, N - dims - 1, dtype)
atg_softmax(ptr, input.ptr, N - dims, dtype)
Tensor{T,N}(ptr[], on(input))
end

function _meanpool(t::Tensor{T,N}, k, s, p, op_sz) where {T,N}
function _meanpool(t::Tensor{T,N}, kernel_size; stride = [1] , padding = [0]) where {T,N}
k = collect(kernel_size)
s = collect(stride)
p = collect(padding)
ptr = Ref(Ptr{Cvoid}())

atg_avg_pool2d(ptr, t.ptr,
k, length(k),
s, length(s),
p, length(p),
0, # ceil_mode
1, # count_include_pad
1 # divisor_override
reverse(k), length(k),
reverse(s), length(s),
reverse(p), length(p),
0, # ceil_mode
1, # count_include_pad
prod(k) # divisor_override
)
Tensor{T,N}(ptr[], on(t))
end

function _maxpool(t::Tensor{T,N}, kernel_size; stride = [1], padding = [0], dilation = [1]) where {T,N}
Copy link
Member

@DhairyaLGandhi DhairyaLGandhi Aug 26, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we need this method?

Copy link
Author

@bjosv79 bjosv79 Aug 26, 2020 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I did not have the intention at first to make a PR, so I changed a few things that was not really needed. Sorry for that.

In my mind it is more clean to have the NNlib ConvDims and PoolDims only in nnlib.jl and not in the ops.jl. Also it makes sense to align the arguments of _meanpool and _maxpool. My own preference would be to remove the _maxpool method with PoolDims.

Copy link
Member

@DhairyaLGandhi DhairyaLGandhi Aug 26, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Best to keep the consistency with the method signatures and API for now.

Copy link
Author

@bjosv79 bjosv79 Sep 1, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was not aware of the #21 issue. Is this PR still interesting? I can clean it up to include only the tests and the changes needed to pass the tests. In other words, remove the added layers, leave the method signatures as they are. Reversing strides, pads, dilations... to match the spatial sizes in ops.jl is however still needed.

I guess this will be resolved by #21 at some point.

Copy link
Member

@DhairyaLGandhi DhairyaLGandhi Sep 1, 2020

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The PR is definitely interesting. #21 should be fixed in the latest release too.

Keeping the signatures consistent is important however. The layers themselves should be fine. Reversing the dims is right too, I think.

k = collect(kernel_size)
s = collect(stride)
p = collect(padding)
d = collect(dilation)
ptr = Ref(Ptr{Cvoid}())

atg_max_pool2d(ptr, t.ptr,
reverse(k), length(k),
reverse(s), length(s),
reverse(p), length(p),
reverse(d), length(d),
0, # ceil_mode
)
Tensor{T,N}(ptr[], on(t))
end
Expand All @@ -129,10 +192,10 @@ function _maxpool(t::Tensor{T,M}, pdims::PoolDims{N,K,S,P,D};
ptr = Ref(Ptr{Cvoid}())

atg_max_pool2d(ptr, t.ptr,
k, length(k),
s, length(s),
p, length(p),
d, length(d),
reverse(k), length(k),
reverse(s), length(s),
reverse(p), length(p),
reverse(d), length(d),
ceil_mode, # ceil_mode
)

Expand All @@ -149,16 +212,67 @@ function _maxpool_with_inds(t::Tensor{T,M}, pdims::PoolDims{N,K,S,P,D};
ptr = [Ptr{Cvoid}(), Ptr{Cvoid}()]

atg_max_pool2d_with_indices(ptr, t.ptr,
k, length(k),
s, length(s),
p, length(p),
d, length(d),
ceil_mode, # ceil_mode
reverse(k), length(k),
reverse(s), length(s),
reverse(p), length(p),
reverse(d), length(d),
ceil_mode,
)

Tensor{T,M}(ptr[1], on(t)), Tensor{T,M}(ptr[2], on(t))
end

function _upsample_nearest2d(t::Tensor{T,N}, output_size) where {T,N}
ptr = Ref(Ptr{Cvoid}())

atg_upsample_nearest2d(ptr, t.ptr,
reverse(output_size), length(output_size),
)
Tensor{T,N}(ptr[], on(t))
end

function _upsample_bilinear2d(t::Tensor{T,N}, output_size, align_corners = true) where {T,N}
ptr = Ref(Ptr{Cvoid}())

atg_upsample_bilinear2d(ptr, t.ptr,
reverse(output_size), length(output_size),
align_corners,
)
Tensor{T,N}(ptr[], on(t))
end

function _upsample_bicubic2d(t::Tensor{T,N}, output_size, align_corners = true) where {T,N}
ptr = Ref(Ptr{Cvoid}())

atg_upsample_bicubic2d(ptr, t.ptr,
reverse(output_size), length(output_size),
align_corners,
)
Tensor{T,N}(ptr[], on(t))
end

function upsample(t::Tensor{T,N}, output_size, mode) where {T,N}
if mode == :NEAREST
_upsample_nearest2d(t, output_size)
elseif mode == :LINEAR
_upsample_bilinear2d(t, output_size)
elseif mode == :CUBIC
_upsample_bicubic2d(t, output_size)
else
error("Unsupported mode $(mode).")
end
end

function pad(t::Tensor{T,N}, padding) where {T,N}
ptr = Ref(Ptr{Cvoid}())
p = collect(padding)

atg_constant_pad_nd(ptr, t.ptr,
p, length(p),
)
Tensor{T,N}(ptr[], on(t))
end

function _chunk(t::Tensor{T,N}, chunks=2, dims=1) where {T,N}
ts = [Ptr{Cvoid}() for _ in 1:chunks]
atg_chunk(ts, t.ptr, chunks, N - dims)
Expand Down
Loading