From 9af22e791e0e8f123882a9091ce5f48f952e819a Mon Sep 17 00:00:00 2001 From: Jesper Stemann Andersen Date: Wed, 8 Jan 2025 08:50:02 +0100 Subject: [PATCH 1/2] Set-up JuliaFormatter --- .JuliaFormatter.toml | 5 +++++ .github/workflows/Format.yml | 8 ++++++++ README.md | 3 +++ deps/julia_wrapper_generator/generator.toml | 2 +- src/Torch.jl | 2 +- src/{wrapper.jl => Wrapper.jl} | 0 6 files changed, 18 insertions(+), 2 deletions(-) create mode 100644 .JuliaFormatter.toml create mode 100644 .github/workflows/Format.yml rename src/{wrapper.jl => Wrapper.jl} (100%) diff --git a/.JuliaFormatter.toml b/.JuliaFormatter.toml new file mode 100644 index 0000000..518dd7f --- /dev/null +++ b/.JuliaFormatter.toml @@ -0,0 +1,5 @@ +style = "blue" + +ignore = ["src/Wrapper.jl"] +whitespace_in_kwargs = true +whitespace_typedefs = true diff --git a/.github/workflows/Format.yml b/.github/workflows/Format.yml new file mode 100644 index 0000000..5a0c29c --- /dev/null +++ b/.github/workflows/Format.yml @@ -0,0 +1,8 @@ +name: Format suggestions +on: + pull_request: +jobs: + code-style: + runs-on: ubuntu-latest + steps: + - uses: julia-actions/julia-format@v3 diff --git a/README.md b/README.md index c0cf008..4795c3b 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,8 @@ # Torch.jl +[![Build Status](https://github.com/FluxML/Torch.jl/actions/workflows/CI.yaml/badge.svg?branch=master)](https://github.com/FluxML/Torch.jl/actions/workflows/CI.yaml?query=branch%3Amaster) +[![Code Style: Blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/JuliaDiff/BlueStyle) + Sensible extensions for exposing torch in Julia. This package is aimed at providing the `Tensor` type, which offloads all computations over to [ATen](https://pytorch.org/cppdocs/), the foundational tensor library for PyTorch, written in C++. diff --git a/deps/julia_wrapper_generator/generator.toml b/deps/julia_wrapper_generator/generator.toml index 1a13a37..d60561e 100644 --- a/deps/julia_wrapper_generator/generator.toml +++ b/deps/julia_wrapper_generator/generator.toml @@ -1,6 +1,6 @@ [general] library_name = "libtorch_c_api" -output_file_path = "../../src/wrapper.jl" +output_file_path = "../../src/Wrapper.jl" prologue_file_path = "./prologue.jl" module_name = "Wrapper" jll_pkg_name = "TorchCAPI_jll" diff --git a/src/Torch.jl b/src/Torch.jl index 3770ca4..5302104 100644 --- a/src/Torch.jl +++ b/src/Torch.jl @@ -13,7 +13,7 @@ using FillArrays TURN_ON_LOGGING = false -include("wrapper.jl") +include("Wrapper.jl") using .Wrapper diff --git a/src/wrapper.jl b/src/Wrapper.jl similarity index 100% rename from src/wrapper.jl rename to src/Wrapper.jl From 9e59d134d0964d79b159d97c93abe37982d56a56 Mon Sep 17 00:00:00 2001 From: Jesper Stemann Andersen Date: Wed, 8 Jan 2025 09:04:31 +0100 Subject: [PATCH 2/2] Implemented formatting --- .JuliaFormatter.toml | 1 + deps/julia_wrapper_generator/generator.jl | 4 +- deps/julia_wrapper_generator/prologue.jl | 20 +- src/Torch.jl | 28 +- src/broadcast.jl | 39 +- src/grads.jl | 294 ++++++++------ src/nnlib.jl | 122 +++--- src/normalise.jl | 145 ++++--- src/ops.jl | 467 +++++++++++++--------- src/scalar.jl | 12 +- src/statistics.jl | 36 +- src/tensor.jl | 214 +++++----- test/flux_tests.jl | 2 +- test/runtests.jl | 2 +- test/tensor_movement_tests.jl | 4 +- test/tensor_nnlib_tests.jl | 132 +++--- 16 files changed, 851 insertions(+), 671 deletions(-) diff --git a/.JuliaFormatter.toml b/.JuliaFormatter.toml index 518dd7f..8a8336e 100644 --- a/.JuliaFormatter.toml +++ b/.JuliaFormatter.toml @@ -1,5 +1,6 @@ style = "blue" ignore = ["src/Wrapper.jl"] +pipe_to_function_call = false whitespace_in_kwargs = true whitespace_typedefs = true diff --git a/deps/julia_wrapper_generator/generator.jl b/deps/julia_wrapper_generator/generator.jl index b24b0b2..9f68499 100644 --- a/deps/julia_wrapper_generator/generator.jl +++ b/deps/julia_wrapper_generator/generator.jl @@ -22,11 +22,11 @@ function rewrite!(e::Expr) end function rewrite!(e::Expr, ::Val{:function}) - rewrite!(e.args[2], Val(e.args[2].head)) + return rewrite!(e.args[2], Val(e.args[2].head)) end function rewrite!(e::Expr, ::Val{:block}) - e.args[1] = Expr(:macrocall, Symbol("@runtime_error_check"), nothing, e.args[1]) + return e.args[1] = Expr(:macrocall, Symbol("@runtime_error_check"), nothing, e.args[1]) end function rewrite!(dag::ExprDAG) diff --git a/deps/julia_wrapper_generator/prologue.jl b/deps/julia_wrapper_generator/prologue.jl index f25db5c..7621fe0 100644 --- a/deps/julia_wrapper_generator/prologue.jl +++ b/deps/julia_wrapper_generator/prologue.jl @@ -1,15 +1,15 @@ function get_error() - err = cglobal((:myerr, libtorch_c_api), Cstring) |> unsafe_load - unsafe_string(err) + err = cglobal((:myerr, libtorch_c_api), Cstring) |> unsafe_load + return unsafe_string(err) end macro runtime_error_check(ex) - quote - x = $ex - if x == 1 - cs = get_error() - flush_error() - throw(cs) - end - end |> esc + return quote + x = $ex + if x == 1 + cs = get_error() + flush_error() + throw(cs) + end + end |> esc end diff --git a/src/Torch.jl b/src/Torch.jl index 5302104..722baae 100644 --- a/src/Torch.jl +++ b/src/Torch.jl @@ -32,23 +32,25 @@ include("statistics.jl") include("grads.jl") include("utils.jl") -@init @require Flux="587475ba-b771-5e3f-ad9e-33799f191a9c" begin - using .Flux +@init @require Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" begin + using .Flux - function (tbn::Flux.BatchNorm)(x::Tensor) - tbn.λ.(Torch.batchnorm(x, tbn.γ, tbn.β, tbn.μ, tbn.σ², 0, tbn.momentum, tbn.ϵ, 1)) - end + function (tbn::Flux.BatchNorm)(x::Tensor) + return tbn.λ.( + Torch.batchnorm(x, tbn.γ, tbn.β, tbn.μ, tbn.σ², 0, tbn.momentum, tbn.ϵ, 1) + ) + end - function Flux.Zygote.accum(t1::Tensor, t2::Tensor{T,N}) where {T,N} - ptr = Ref(Ptr{Cvoid}()) + function Flux.Zygote.accum(t1::Tensor, t2::Tensor{T, N}) where {T, N} + ptr = Ref(Ptr{Cvoid}()) - Torch.Wrapper.atg_add_(ptr, t1.ptr, t2.ptr) - Tensor{T,N}(ptr[], Torch.on(t1)) - end + Torch.Wrapper.atg_add_(ptr, t1.ptr, t2.ptr) + return Tensor{T, N}(ptr[], Torch.on(t1)) + end - eval(:(Flux.Zygote.@nograd Torch.Wrapper.at_copy_data)) - eval(:(Flux.Zygote.@nograd Torch.Wrapper.at_dim)) - torch(x) = Flux.fmap(to_tensor, x) + eval(:(Flux.Zygote.@nograd Torch.Wrapper.at_copy_data)) + eval(:(Flux.Zygote.@nograd Torch.Wrapper.at_dim)) + torch(x) = Flux.fmap(to_tensor, x) end end # module diff --git a/src/broadcast.jl b/src/broadcast.jl index 7d6499b..d6cbdfb 100644 --- a/src/broadcast.jl +++ b/src/broadcast.jl @@ -9,23 +9,23 @@ using Base.Broadcast: broadcast_shape # Base.BroadcastStyle(::Type{Tensor}) = TensorStyle() for op in (:+, :-, :/) - @eval function broadcasted(::typeof($op), t1::Tensor, t2::Tensor) - $op(t1, t2) - end + @eval function broadcasted(::typeof($op), t1::Tensor, t2::Tensor) + return $op(t1, t2) + end end for op in (:+, :-) - @eval function broadcasted(::typeof($op), t1::Tensor, t2::TensorVector) - t_ = reshape(t2, -1, 1) - $op(t1, t_) - end + @eval function broadcasted(::typeof($op), t1::Tensor, t2::TensorVector) + t_ = reshape(t2, -1, 1) + return $op(t1, t_) + end end -function broadcasted(::typeof(*), t1::Tensor{T,N}, t2::Tensor{T,M}) where {T,N,M} - ptr = Ref(Ptr{Cvoid}()) +function broadcasted(::typeof(*), t1::Tensor{T, N}, t2::Tensor{T, M}) where {T, N, M} + ptr = Ref(Ptr{Cvoid}()) - atg_mul(ptr, t1.ptr, t2.ptr) - Tensor{T,max(N,M)}(ptr[], on(t1)) + atg_mul(ptr, t1.ptr, t2.ptr) + return Tensor{T, max(N, M)}(ptr[], on(t1)) end broadcasted(::typeof(NNlib.relu), t::Tensor) = NNlib.relu(t) @@ -34,22 +34,21 @@ broadcasted(::typeof(identity), t::Tensor) = identity(t) broadcasted(::typeof(NNlib.sigmoid), t::Tensor) = NNlib.sigmoid(t) for op in (:+, :-, :*, :/) - @eval function broadcasted(::typeof($op), t::Tensor, args...) - $op(t, args...) - end + @eval function broadcasted(::typeof($op), t::Tensor, args...) + return $op(t, args...) + end end broadcasted(::typeof(sqrt), t::Tensor) = sqrt(t) -function broadcasted(::typeof(copy), t::Tensor{T,N}) where {T,N} - t +function broadcasted(::typeof(copy), t::Tensor{T, N}) where {T, N} + return t end @adjoint function broadcast(::typeof(NNlib.sigmoid), t::Tensor) - - NNlib.sigmoid(t), Δ -> (∇sigmoid(Δ, t),) + return NNlib.sigmoid(t), Δ -> (∇sigmoid(Δ, t),) end -@adjoint function broadcasted(::typeof(NNlib.relu), t::Tensor{T}) where T - relu(t), Δ -> (nothing, ∇leaky_relu(Δ, t, zero(T)),) +@adjoint function broadcasted(::typeof(NNlib.relu), t::Tensor{T}) where {T} + return relu(t), Δ -> (nothing, ∇leaky_relu(Δ, t, zero(T))) end diff --git a/src/grads.jl b/src/grads.jl index ff2aae1..dabe363 100644 --- a/src/grads.jl +++ b/src/grads.jl @@ -1,153 +1,201 @@ import NNlib: ∇conv_data, ∇conv_filter @adjoint function tensor(x; kwargs...) - tensor(x; kwargs...), Δ -> (collect(Δ),) + return tensor(x; kwargs...), Δ -> (collect(Δ),) end -function cudnn_convolution_backward_bias(t::Tensor{T,N}) where {T,N} - ptr = Ref(Ptr{Cvoid}()) - atg_cudnn_convolution_backward_bias(ptr, t.ptr) - Tensor{T,N}(ptr[], on(t)) +function cudnn_convolution_backward_bias(t::Tensor{T, N}) where {T, N} + ptr = Ref(Ptr{Cvoid}()) + atg_cudnn_convolution_backward_bias(ptr, t.ptr) + return Tensor{T, N}(ptr[], on(t)) end const ∇conv_bias = cudnn_convolution_backward_bias -function ∇conv_data(dy::AbstractArray, w::Tensor{T}, - cdims::DenseConvDims{M,K,S,P,D}; - groups = 1, - benchmark = 0, - deterministic = 0, - allow_tf32 = 0) where {M,K,S,P,D,T} - - ptr = Ref(Ptr{Cvoid}()) - dy_ = tensor(dy, dev = on(w)) - padding = NNlib.padding(cdims) - padding = [padding[1];padding[3]] - stride = collect(NNlib.stride(cdims)) - dilation = collect(NNlib.dilation(cdims)) - - s = reverse([NNlib.input_size(cdims)..., - NNlib.channels_in(cdims), - size(dy_, ndims(dy_))]) - - atg_cudnn_convolution_backward_input(ptr, - s, length(s), - dy_.ptr, w.ptr, - padding, length(padding), - stride, length(stride), - dilation, length(dilation), - groups, benchmark, deterministic, allow_tf32) - Tensor{T,ndims(dy_)}(ptr[], on(dy_)) +function ∇conv_data( + dy::AbstractArray, + w::Tensor{T}, + cdims::DenseConvDims{M, K, S, P, D}; + groups = 1, + benchmark = 0, + deterministic = 0, + allow_tf32 = 0, +) where {M, K, S, P, D, T} + ptr = Ref(Ptr{Cvoid}()) + dy_ = tensor(dy; dev = on(w)) + padding = NNlib.padding(cdims) + padding = [padding[1]; padding[3]] + stride = collect(NNlib.stride(cdims)) + dilation = collect(NNlib.dilation(cdims)) + + s = reverse([ + NNlib.input_size(cdims)..., NNlib.channels_in(cdims), size(dy_, ndims(dy_)) + ]) + + atg_cudnn_convolution_backward_input( + ptr, + s, + length(s), + dy_.ptr, + w.ptr, + padding, + length(padding), + stride, + length(stride), + dilation, + length(dilation), + groups, + benchmark, + deterministic, + allow_tf32, + ) + return Tensor{T, ndims(dy_)}(ptr[], on(dy_)) end -function ∇conv_filter(w::Tensor{T}, dy::AbstractArray{T}, - cdims::DenseConvDims{M,K,S,P,D}; - groups = 1, - benchmark = 0, - deterministic = 0, - allow_tf32 = 0) where {M,K,S,P,D,T} - - dy_ = tensor(dy, dev = on(w)) - ptr = Ref(Ptr{Cvoid}()) - padding = NNlib.padding(cdims) - padding = [padding[1];padding[3]] - stride = collect(NNlib.stride(cdims)) - dilation = collect(NNlib.dilation(cdims)) - - s = reverse([NNlib.kernel_size(cdims)..., - NNlib.channels_in(cdims), - NNlib.channels_out(cdims)]) - - atg_cudnn_convolution_backward_weight(ptr, - s, length(s), - dy_.ptr, w.ptr, - padding, length(padding), - stride, length(stride), - dilation, length(dilation), - groups, benchmark, deterministic, allow_tf32) - - Tensor{T,ndims(dy_)}(ptr[], on(dy_)) +function ∇conv_filter( + w::Tensor{T}, + dy::AbstractArray{T}, + cdims::DenseConvDims{M, K, S, P, D}; + groups = 1, + benchmark = 0, + deterministic = 0, + allow_tf32 = 0, +) where {M, K, S, P, D, T} + dy_ = tensor(dy; dev = on(w)) + ptr = Ref(Ptr{Cvoid}()) + padding = NNlib.padding(cdims) + padding = [padding[1]; padding[3]] + stride = collect(NNlib.stride(cdims)) + dilation = collect(NNlib.dilation(cdims)) + + s = reverse([ + NNlib.kernel_size(cdims)..., NNlib.channels_in(cdims), NNlib.channels_out(cdims) + ]) + + atg_cudnn_convolution_backward_weight( + ptr, + s, + length(s), + dy_.ptr, + w.ptr, + padding, + length(padding), + stride, + length(stride), + dilation, + length(dilation), + groups, + benchmark, + deterministic, + allow_tf32, + ) + + return Tensor{T, ndims(dy_)}(ptr[], on(dy_)) end -function NNlib.∇maxpool(dy::AbstractArray{T}, y::Tensor{T,M}, x::Tensor{T,M}, - pdims::PoolDims{N,K,S,P,D}; - ceil_mode = 0, - indices=nothing) where {N,K,S,P,D, T,M} - - dy_ = tensor(dy, dev = on(y)) - ptr = Ref(Ptr{Cvoid}()) - kernel = collect(NNlib.kernel_size(pdims)) - stride = collect(NNlib.stride(pdims)) - padding = NNlib.padding(pdims) - padding = Int[padding[1];padding[3]] - dilation = collect(NNlib.dilation(pdims)) - atg_max_pool2d_with_indices_backward(ptr, dy_.ptr, x.ptr, - kernel, length(kernel), - stride, length(stride), - padding, length(padding), - dilation, length(dilation), - ceil_mode, - indices.ptr - ) - - mp = Tensor{T,N}(ptr[], on(x)) - reshape(mp, reverse(size(mp))...) +function NNlib.∇maxpool( + dy::AbstractArray{T}, + y::Tensor{T, M}, + x::Tensor{T, M}, + pdims::PoolDims{N, K, S, P, D}; + ceil_mode = 0, + indices = nothing, +) where {N, K, S, P, D, T, M} + dy_ = tensor(dy; dev = on(y)) + ptr = Ref(Ptr{Cvoid}()) + kernel = collect(NNlib.kernel_size(pdims)) + stride = collect(NNlib.stride(pdims)) + padding = NNlib.padding(pdims) + padding = Int[padding[1]; padding[3]] + dilation = collect(NNlib.dilation(pdims)) + atg_max_pool2d_with_indices_backward( + ptr, + dy_.ptr, + x.ptr, + kernel, + length(kernel), + stride, + length(stride), + padding, + length(padding), + dilation, + length(dilation), + ceil_mode, + indices.ptr, + ) + + mp = Tensor{T, N}(ptr[], on(x)) + return reshape(mp, reverse(size(mp))...) end @adjoint function NNlib.maxpool(t::Tensor, pdims::PoolDims; ceil_mode = 0) - y, inds = _maxpool_with_inds(t, pdims, ceil_mode = ceil_mode) - y, Δ -> begin - (∇maxpool(Δ, y, t, pdims, ceil_mode = ceil_mode, indices = inds), nothing) - end + y, inds = _maxpool_with_inds(t, pdims; ceil_mode = ceil_mode) + return y, + Δ -> begin + (∇maxpool(Δ, y, t, pdims; ceil_mode = ceil_mode, indices = inds), nothing) + end end -function NNlib.∇meanpool(dy::AbstractArray, y::Tensor{T,M}, x::Tensor{T,M}, - pdims::PoolDims{N,K,S,P,D}; - ceil_mode = 0, - count_include_pad = 1, - divisor_override = 1) where {N,K,S,P,D, T,M} - - ptr = Ref(Ptr{Cvoid}()) - dy_ = dy isa Base.ReshapedArray ? reshape(parent(dy), dy.dims...) : tensor(dy, dev = on(y)) - kernel = collect(NNlib.kernel_size(pdims)) - stride = collect(NNlib.stride(pdims)) - padding = NNlib.padding(pdims) - padding = [padding[1];padding[3]] - - atg_avg_pool2d_backward(ptr, - dy_.ptr, x.ptr, - kernel, length(kernel), - stride, length(stride), - padding, length(padding), - ceil_mode, - count_include_pad, - divisor_override) - - Tensor{T,M}(ptr[], on(x)) +function NNlib.∇meanpool( + dy::AbstractArray, + y::Tensor{T, M}, + x::Tensor{T, M}, + pdims::PoolDims{N, K, S, P, D}; + ceil_mode = 0, + count_include_pad = 1, + divisor_override = 1, +) where {N, K, S, P, D, T, M} + ptr = Ref(Ptr{Cvoid}()) + dy_ = if dy isa Base.ReshapedArray + reshape(parent(dy), dy.dims...) + else + tensor(dy; dev = on(y)) + end + kernel = collect(NNlib.kernel_size(pdims)) + stride = collect(NNlib.stride(pdims)) + padding = NNlib.padding(pdims) + padding = [padding[1]; padding[3]] + + atg_avg_pool2d_backward( + ptr, + dy_.ptr, + x.ptr, + kernel, + length(kernel), + stride, + length(stride), + padding, + length(padding), + ceil_mode, + count_include_pad, + divisor_override, + ) + + return Tensor{T, M}(ptr[], on(x)) end -function ∇sigmoid(dy::AbstractArray, t::Tensor{T,N}) where {T,N} - ptr = Ref(Ptr{Cvoid}()) +function ∇sigmoid(dy::AbstractArray, t::Tensor{T, N}) where {T, N} + ptr = Ref(Ptr{Cvoid}()) - dy_ = tensor(dy, dev = on(t)) - atg_sigmoid_backward(ptr, dy_.ptr, t.ptr) - Tensor{T,N}(ptr[], on(t)) + dy_ = tensor(dy; dev = on(t)) + atg_sigmoid_backward(ptr, dy_.ptr, t.ptr) + return Tensor{T, N}(ptr[], on(t)) end @adjoint function NNlib.sigmoid(t::Tensor) - x = sigmoid(t) - x, Δ -> (∇sigmoid(Δ, x),) + x = sigmoid(t) + return x, Δ -> (∇sigmoid(Δ, x),) end -function ∇leaky_relu(dy::AbstractArray, x::Tensor{T,N}, slope) where {T,N} - ptr = Ref(Ptr{Cvoid}()) +function ∇leaky_relu(dy::AbstractArray, x::Tensor{T, N}, slope) where {T, N} + ptr = Ref(Ptr{Cvoid}()) - dy_ = tensor(dy, dev = on(x)) - self_is_result = 0 - atg_leaky_relu_backward(ptr, dy_.ptr, x.ptr, Scalar(slope).ptr, self_is_result) - Tensor{T,N}(ptr[], on(x)) + dy_ = tensor(dy; dev = on(x)) + self_is_result = 0 + atg_leaky_relu_backward(ptr, dy_.ptr, x.ptr, Scalar(slope).ptr, self_is_result) + return Tensor{T, N}(ptr[], on(x)) end -@adjoint function NNlib.relu(t::Tensor{T,N}) where {T,N} - NNlib.relu(t), Δ -> (∇leaky_relu(Δ, t, zero(T)),) +@adjoint function NNlib.relu(t::Tensor{T, N}) where {T, N} + return NNlib.relu(t), Δ -> (∇leaky_relu(Δ, t, zero(T)),) end diff --git a/src/nnlib.jl b/src/nnlib.jl index 07c18ef..c93c26c 100644 --- a/src/nnlib.jl +++ b/src/nnlib.jl @@ -4,87 +4,105 @@ using NNlib: PoolDims import NNlib: conv, depthwiseconv -function NNlib.conv(x::Tensor{xT, N}, w::Tensor, b::Tensor{T}, - cdims::DenseConvDims{M,K,S,P,D}) where {T,N,xT,M,K,S,P,D} - stride = NNlib.stride(cdims) - padding = NNlib.padding(cdims) - dilation = NNlib.dilation(cdims) - op = conv2d(x, w, b, stride = collect(stride), padding = [padding[1];padding[3]], dilation = collect(dilation)) - op +function NNlib.conv( + x::Tensor{xT, N}, w::Tensor, b::Tensor{T}, cdims::DenseConvDims{M, K, S, P, D} +) where {T, N, xT, M, K, S, P, D} + stride = NNlib.stride(cdims) + padding = NNlib.padding(cdims) + dilation = NNlib.dilation(cdims) + op = conv2d( + x, + w, + b; + stride = collect(stride), + padding = [padding[1]; padding[3]], + dilation = collect(dilation), + ) + return op end 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) - op + b = zeros(Tensor{Float32}, size(w)[end]; dev = on(w)) + op = conv(x, w, b, cdims) + return 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 +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), + ) + return 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 + b = zeros(Tensor{Float32}, size(w)[end]; dev = on(w)) + op = depthwiseconv( + x, w, b; stride = collect(stride), pad = collect(pad), dilation = collect(dilation) + ) + return op end -function NNlib.relu(t::Tensor{T,N}) where {T,N} - ptr = Ref(Ptr{Cvoid}()) +function NNlib.relu(t::Tensor{T, N}) where {T, N} + ptr = Ref(Ptr{Cvoid}()) - atg_relu(ptr, t.ptr) - Tensor{T,N}(ptr[], on(t)) + atg_relu(ptr, t.ptr) + return Tensor{T, N}(ptr[], on(t)) end -function NNlib.leakyrelu(t::Tensor{T,N}) where {T,N} - ptr = Ref(Ptr{Cvoid}()) +function NNlib.leakyrelu(t::Tensor{T, N}) where {T, N} + ptr = Ref(Ptr{Cvoid}()) - atg_leaky_relu(ptr, t.ptr) - Tensor{T,N}(ptr[], on(t)) + atg_leaky_relu(ptr, t.ptr) + return Tensor{T, N}(ptr[], on(t)) end -function NNlib.sigmoid(t::Tensor{T,N}) where {T,N} - ptr = Ref(Ptr{Cvoid}()) +function NNlib.sigmoid(t::Tensor{T, N}) where {T, N} + ptr = Ref(Ptr{Cvoid}()) - atg_sigmoid(ptr, t.ptr) - Tensor{T,N}(ptr[], on(t)) + atg_sigmoid(ptr, t.ptr) + return Tensor{T, N}(ptr[], on(t)) end -function NNlib.tanh(t::Tensor{T,N}) where {T,N} - ptr = Ref(Ptr{Cvoid}()) +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)) + atg_tanh(ptr, t.ptr) + return Tensor{T, N}(ptr[], on(t)) end -function NNlib.softmax(t::Tensor{T,N}; dims = 1) where {T,N} - _softmax(t, dims, options[T]) +function NNlib.softmax(t::Tensor{T, N}; dims = 1) where {T, N} + return _softmax(t, dims, options[T]) end function NNlib.∇softmax(Δ, xs::Tensor; dims = 1) - t = tensor(Δ, dev = on(xs)) - sf = softmax(xs, dims=dims) - sf .* (t .- sum(t .* sf, dims = dims)) + t = tensor(Δ; dev = on(xs)) + sf = softmax(xs; dims = dims) + return sf .* (t .- sum(t .* sf; dims = dims)) end -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(NNlib.stride(pdims)) - padding = NNlib.padding(pdims) - # op_sz = NNlib.output_size(pdims) +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(NNlib.stride(pdims)) + padding = NNlib.padding(pdims) + # op_sz = NNlib.output_size(pdims) - _meanpool(t, ks, stride=stride, padding=[padding[1];padding[3]]) + return _meanpool(t, ks; stride = stride, padding = [padding[1]; padding[3]]) end -function NNlib.maxpool(t::Tensor, pdims::PoolDims{N,K,S,P,D}) where {N,K,S,P,D} - ks = collect(NNlib.kernel_size(pdims)) - stride = collect(NNlib.stride(pdims)) - padding = NNlib.padding(pdims) - dilation = collect(NNlib.dilation(pdims)) +function NNlib.maxpool(t::Tensor, pdims::PoolDims{N, K, S, P, D}) where {N, K, S, P, D} + ks = collect(NNlib.kernel_size(pdims)) + stride = collect(NNlib.stride(pdims)) + padding = NNlib.padding(pdims) + dilation = collect(NNlib.dilation(pdims)) - _maxpool(t, ks, stride=stride, padding=[padding[1];padding[3]], dilation=dilation) + return _maxpool( + t, ks; stride = stride, padding = [padding[1]; padding[3]], dilation = dilation + ) end diff --git a/src/normalise.jl b/src/normalise.jl index 18c4e95..96b95da 100644 --- a/src/normalise.jl +++ b/src/normalise.jl @@ -1,61 +1,106 @@ -function batchnorm(input::Tensor{T,N}, weight, bias, - running_mean, running_var, - training, momentum, - ep, cudnn_enabled) where {T,N} - - ptr = Ref(Ptr{Cvoid}()) - atg_batch_norm(ptr, input.ptr, - weight.ptr, bias.ptr, - running_mean.ptr, running_var.ptr, - training, momentum, - ep, cudnn_enabled) - - Tensor{T,N}(ptr[], on(input)) -end +function batchnorm( + input::Tensor{T, N}, + weight, + bias, + running_mean, + running_var, + training, + momentum, + ep, + cudnn_enabled, +) where {T, N} + ptr = Ref(Ptr{Cvoid}()) + atg_batch_norm( + ptr, + input.ptr, + weight.ptr, + bias.ptr, + running_mean.ptr, + running_var.ptr, + training, + momentum, + ep, + cudnn_enabled, + ) -function ∇batchnorm(dy::AbstractArray, input::Tensor{T,N}, - running_mean, invstd, weight::Tensor{T,S}, - input_g = 1, weight_g = 1, bias_g = 1) where {T,N,S} + return Tensor{T, N}(ptr[], on(input)) +end - ptr = [Ptr{Cvoid}() for i = 1:4] - dy_ = tensor(dy, dev = on(input)) +function ∇batchnorm( + dy::AbstractArray, + input::Tensor{T, N}, + running_mean, + invstd, + weight::Tensor{T, S}, + input_g = 1, + weight_g = 1, + bias_g = 1, +) where {T, N, S} + ptr = [Ptr{Cvoid}() for i in 1:4] + dy_ = tensor(dy; dev = on(input)) - atg_batch_norm_backward_reduce(ptr, dy_.ptr, input.ptr, - running_mean.ptr, invstd.ptr, - weight.ptr, - input_g, weight_g, bias_g) + atg_batch_norm_backward_reduce( + ptr, + dy_.ptr, + input.ptr, + running_mean.ptr, + invstd.ptr, + weight.ptr, + input_g, + weight_g, + bias_g, + ) - rank = Ref{Cint}(-1) - [Tensor{T,S}(i, on(input)) for i in ptr[1:4]] + rank = Ref{Cint}(-1) + return [Tensor{T, S}(i, on(input)) for i in ptr[1:4]] end +function ∇batchnorm_element( + dy::AbstractArray, + input::Tensor{T, N}, + running_mean, + invstd, + weight::Tensor{T, S}, + mean_dy = nothing, + mean_dy_xmu = nothing, +) where {T, N, S} + ptr = Ref(Ptr{Cvoid}()) + dy_ = tensor(dy; dev = on(input)) + count = Tensor(Int32, size(input)...; dev = on(input)) -function ∇batchnorm_element(dy::AbstractArray, input::Tensor{T,N}, - running_mean, invstd, weight::Tensor{T,S}, - mean_dy = nothing, mean_dy_xmu = nothing) where {T,N,S} - - ptr = Ref(Ptr{Cvoid}()) - dy_ = tensor(dy, dev = on(input)) - count = Tensor(Int32, size(input)..., dev = on(input)) + atg_batch_norm_backward_elemt( + ptr, + dy_.ptr, + input.ptr, + running_mean.ptr, + invstd.ptr, + weight.ptr, + weight.ptr, + weight.ptr, + count.ptr, + ) # Hack passing weight.ptr as mean_dy, and mean_dy_xmu - atg_batch_norm_backward_elemt(ptr, dy_.ptr, input.ptr, - running_mean.ptr, invstd.ptr, - weight.ptr, - weight.ptr, weight.ptr, count.ptr) # Hack passing weight.ptr as mean_dy, and mean_dy_xmu - - Tensor{T,N}(ptr[], on(input)) + return Tensor{T, N}(ptr[], on(input)) end -@adjoint function batchnorm(input, weight, bias, - running_mean, running_var, - training, momentum, - ep, cudnn_enabled) where {T,N} - - y = batchnorm(input, weight, bias, running_mean, running_var, training, momentum, ep, cudnn_enabled) +@adjoint function batchnorm( + input, weight, bias, running_mean, running_var, training, momentum, ep, cudnn_enabled +) where {T, N} + y = batchnorm( + input, + weight, + bias, + running_mean, + running_var, + training, + momentum, + ep, + cudnn_enabled, + ) - y, Δ -> begin - e = ∇batchnorm_element(Δ, input, running_mean, running_var, weight) - vs = ∇batchnorm(Δ, input, running_mean, running_var, weight) - (e, vs..., nothing, nothing, nothing, nothing) - end -end + return y, Δ -> begin + e = ∇batchnorm_element(Δ, input, running_mean, running_var, weight) + vs = ∇batchnorm(Δ, input, running_mean, running_var, weight) + (e, vs..., nothing, nothing, nothing, nothing) + end +end diff --git a/src/ops.jl b/src/ops.jl index 62b1979..380a290 100644 --- a/src/ops.jl +++ b/src/ops.jl @@ -1,126 +1,160 @@ import Base: +, -, *, / using LinearAlgebra -for (op,fn) in zip((:+, :-, :/, :*), (atg_add, atg_sub, atg_div, atg_matmul)) - @eval function $op(t1::Tensor{T,N}, t2::Tensor{T,K}) where {T,N,K} - ptr = Ref(Ptr{Cvoid}()) - rank = Ref{Cint}(-1) +for (op, fn) in zip((:+, :-, :/, :*), (atg_add, atg_sub, atg_div, atg_matmul)) + @eval function $op(t1::Tensor{T, N}, t2::Tensor{T, K}) where {T, N, K} + ptr = Ref(Ptr{Cvoid}()) + rank = Ref{Cint}(-1) - $fn(ptr, t1.ptr, t2.ptr) + $fn(ptr, t1.ptr, t2.ptr) - # TODO: using `rank` here causes compiler to emit error - # make shape checking more robust - # at_dim(rank, ptr[]) - Tensor{T,max(N,K)}(ptr[], on(t1)) - end + # TODO: using `rank` here causes compiler to emit error + # make shape checking more robust + # at_dim(rank, ptr[]) + return Tensor{T, max(N, K)}(ptr[], on(t1)) + end end for op in (:+, :-, :/, :*) - @eval function $op(t::Tensor{T,N}, r::S) where {T,N,S <: Real} - i = T[r] - t2 = tensor(i, dev = on(t)) - res = $op(t, t2) - res - end + @eval function $op(t::Tensor{T, N}, r::S) where {T, N, S <: Real} + i = T[r] + t2 = tensor(i; dev = on(t)) + res = $op(t, t2) + return res + end end # Basic LinAlg handling function LinearAlgebra.adjoint!(dest::Tensor, src::TensorMatrix{T}) where {T} - ptr = Ref(Ptr{Cvoid}()) - p = parent(src) - atg_t(ptr, p.ptr) - t = TensorMatrix{T}(ptr[], on(p)) - at_copy_(dest.ptr, t.ptr) + ptr = Ref(Ptr{Cvoid}()) + p = parent(src) + atg_t(ptr, p.ptr) + t = TensorMatrix{T}(ptr[], on(p)) + return at_copy_(dest.ptr, t.ptr) end -function Base.copyto!(dest::Array, src::Union{Adjoint{T, <:TensorMatrix{T}}, - Transpose{T, <:TensorMatrix{T}}}) where T - t = tensor(similar(src), dev = on(parent(src))) - LinearAlgebra.adjoint!(t, src) - copyto!(dest, t) +function Base.copyto!( + dest::Array, src::Union{Adjoint{T, <:TensorMatrix{T}}, Transpose{T, <:TensorMatrix{T}}} +) where {T} + t = tensor(similar(src); dev = on(parent(src))) + LinearAlgebra.adjoint!(t, src) + return copyto!(dest, t) end -*(a::Fill, b::Tensor) = tensor(a,dev = on(b)) * b -*(a::Fill, b::LinearAlgebra.Adjoint{T,<:Tensor{T}}) where T = tensor(a,dev = on(parent(b))) * b -*(a::LinearAlgebra.Adjoint{T,<:Tensor{T}}, b::Fill) where T = a * tensor(b,dev = on(parent(a))) +*(a::Fill, b::Tensor) = tensor(a; dev = on(b)) * b +function *(a::Fill, b::LinearAlgebra.Adjoint{T, <:Tensor{T}}) where {T} + return tensor(a; dev = on(parent(b))) * b +end +function *(a::LinearAlgebra.Adjoint{T, <:Tensor{T}}, b::Fill) where {T} + return a * tensor(b; dev = on(parent(a))) +end -function *(a::Tensor, b::Union{Adjoint{T, <:TensorMatrix{T}}, - Transpose{T, <:TensorMatrix{T}}}) where T - ptr = Ref(Ptr{Cvoid}()) - p = Tensor(size(b)..., dev = on(a)) - LinearAlgebra.adjoint!(p, b.parent) - a * p +function *( + a::Tensor, b::Union{Adjoint{T, <:TensorMatrix{T}}, Transpose{T, <:TensorMatrix{T}}} +) where {T} + ptr = Ref(Ptr{Cvoid}()) + p = Tensor(size(b)...; dev = on(a)) + LinearAlgebra.adjoint!(p, b.parent) + return a * p end -function *(a::Union{Adjoint{T, <:TensorMatrix{T}}, - Transpose{T, <:TensorMatrix{T}}}, b::Tensor) where T - ptr = Ref(Ptr{Cvoid}()) - p = Tensor(size(a)..., dev = on(b)) - LinearAlgebra.adjoint!(p, a.parent) - p * b +function *( + a::Union{Adjoint{T, <:TensorMatrix{T}}, Transpose{T, <:TensorMatrix{T}}}, b::Tensor +) where {T} + ptr = Ref(Ptr{Cvoid}()) + p = Tensor(size(a)...; dev = on(b)) + LinearAlgebra.adjoint!(p, a.parent) + return p * b end -function Base.maximum(t::Tensor{T}; dims = :) where T - ptr = Ref(Ptr{Cvoid}()) - atg_max(ptr, t.ptr) - Tensor{T,0}(ptr[], on(t)) +function Base.maximum(t::Tensor{T}; dims = :) where {T} + ptr = Ref(Ptr{Cvoid}()) + atg_max(ptr, t.ptr) + return Tensor{T, 0}(ptr[], on(t)) end -function Base.sqrt(t::Tensor{T,N}) where {T,N} - ptr = Ref(Ptr{Cvoid}()) - atg_sqrt(ptr, t.ptr) - Tensor{T,N}(ptr[], on(t)) +function Base.sqrt(t::Tensor{T, N}) where {T, N} + ptr = Ref(Ptr{Cvoid}()) + atg_sqrt(ptr, t.ptr) + return Tensor{T, N}(ptr[], on(t)) end -function Base.cat(ts::Tensor{T,N}...; dims = 1) where {T,N} - ptr = Ref(Ptr{Cvoid}()) - ts_arr = [i.ptr for i in ts] - atg_cat(ptr, ts_arr, length(ts_arr), N - dims) - Tensor{T,N}(ptr[], on(ts[1])) +function Base.cat(ts::Tensor{T, N}...; dims = 1) where {T, N} + ptr = Ref(Ptr{Cvoid}()) + ts_arr = [i.ptr for i in ts] + atg_cat(ptr, ts_arr, length(ts_arr), N - dims) + return Tensor{T, N}(ptr[], on(ts[1])) 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} - - 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) +function conv2d( + input::Tensor{T}, + filter::Tensor{T, N}, + bias::Tensor{T}; + stride = [1], + padding = [0], + dilation = [1], + groups = 1, +) where {T, N} + ptr = Ref(Ptr{Cvoid}()) - Tensor{T,N}(ptr[], on(input)) + atg_conv2d( + ptr, + input.ptr, + filter.ptr, + bias.ptr, + reverse(stride), + length(stride), + reverse(padding), + length(padding), + reverse(dilation), + length(dilation), + groups, + ) + + return 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)) +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}()) - Tensor{T,N}(ptr[], on(input)) + 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), + ) + + return 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} +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. @@ -132,128 +166,163 @@ function _depthwise_conv2d(input::Tensor{T}, filter::Tensor{T,N}, bias::Tensor{T 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)) + atg_conv2d( + ptr, + input.ptr, + filter.ptr, + bias.ptr, + reverse(stride), + length(stride), + reverse(padding), + length(padding), + reverse(dilation), + length(dilation), + groups, + ) + + return 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}()) +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, dtype) - Tensor{T,N}(ptr[], on(input)) + atg_softmax(ptr, input.ptr, N - dims, dtype) + return Tensor{T, N}(ptr[], on(input)) end -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, - 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)) +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, + reverse(k), + length(k), + reverse(s), + length(s), + reverse(p), + length(p), + 0, # ceil_mode + 1, # count_include_pad + prod(k), # divisor_override + ) + return Tensor{T, N}(ptr[], on(t)) end -function _maxpool(t::Tensor{T,N}, kernel_size; stride = [1], padding = [0], dilation = [1]) where {T,N} - 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)) +function _maxpool( + t::Tensor{T, N}, kernel_size; stride = [1], padding = [0], dilation = [1] +) where {T, N} + 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 + ) + return Tensor{T, N}(ptr[], on(t)) end -function _maxpool(t::Tensor{T,M}, pdims::PoolDims{N,K,S,P,D}; - ceil_mode = 0) where {N,K,S,P,D, T,M} - k = collect(NNlib.kernel_size(pdims)) - s = collect(NNlib.stride(pdims)) - padding = NNlib.padding(pdims) - p = [padding[1];padding[3]] - d = collect(NNlib.dilation(pdims)) - - 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), - ceil_mode, # ceil_mode - ) - - Tensor{T,M}(ptr[], on(t)) +function _maxpool( + t::Tensor{T, M}, pdims::PoolDims{N, K, S, P, D}; ceil_mode = 0 +) where {N, K, S, P, D, T, M} + k = collect(NNlib.kernel_size(pdims)) + s = collect(NNlib.stride(pdims)) + padding = NNlib.padding(pdims) + p = [padding[1]; padding[3]] + d = collect(NNlib.dilation(pdims)) + + 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), + ceil_mode, # ceil_mode + ) + + return Tensor{T, M}(ptr[], on(t)) end -function _maxpool_with_inds(t::Tensor{T,M}, pdims::PoolDims{N,K,S,P,D}; - ceil_mode = 0) where {N,K,S,P,D, T,M} - k = collect(NNlib.kernel_size(pdims)) - s = collect(NNlib.stride(pdims)) - padding = NNlib.padding(pdims) - p = [padding[1];padding[3]] - d = collect(NNlib.dilation(pdims)) - - ptr = [Ptr{Cvoid}(), Ptr{Cvoid}()] - - atg_max_pool2d_with_indices(ptr, t.ptr, - 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)) +function _maxpool_with_inds( + t::Tensor{T, M}, pdims::PoolDims{N, K, S, P, D}; ceil_mode = 0 +) where {N, K, S, P, D, T, M} + k = collect(NNlib.kernel_size(pdims)) + s = collect(NNlib.stride(pdims)) + padding = NNlib.padding(pdims) + p = [padding[1]; padding[3]] + d = collect(NNlib.dilation(pdims)) + + ptr = [Ptr{Cvoid}(), Ptr{Cvoid}()] + + atg_max_pool2d_with_indices( + ptr, + t.ptr, + reverse(k), + length(k), + reverse(s), + length(s), + reverse(p), + length(p), + reverse(d), + length(d), + ceil_mode, + ) + + return 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}()) +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)) + atg_upsample_nearest2d(ptr, t.ptr, reverse(output_size), length(output_size)) + return 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}()) +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)) + atg_upsample_bilinear2d( + ptr, t.ptr, reverse(output_size), length(output_size), align_corners + ) + return 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}()) +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)) + atg_upsample_bicubic2d( + ptr, t.ptr, reverse(output_size), length(output_size), align_corners + ) + return Tensor{T, N}(ptr[], on(t)) end -function upsample(t::Tensor{T,N}, output_size, mode) where {T,N} +function upsample(t::Tensor{T, N}, output_size, mode) where {T, N} if mode == :NEAREST _upsample_nearest2d(t, output_size) elseif mode == :LINEAR @@ -261,22 +330,20 @@ function upsample(t::Tensor{T,N}, output_size, mode) where {T,N} elseif mode == :CUBIC _upsample_bicubic2d(t, output_size) else - error("Unsupported mode $(mode).") + error("Unsupported mode $(mode).") end end -function pad(t::Tensor{T,N}, padding) where {T,N} - ptr = Ref(Ptr{Cvoid}()) - p = collect(padding) +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)) + atg_constant_pad_nd(ptr, t.ptr, p, length(p)) + return 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) - [Tensor{T,N}(ts[i], on(t)) for i in 1:chunks] +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) + return [Tensor{T, N}(ts[i], on(t)) for i in 1:chunks] end diff --git a/src/scalar.jl b/src/scalar.jl index c6c21ed..be5d2de 100644 --- a/src/scalar.jl +++ b/src/scalar.jl @@ -1,13 +1,13 @@ mutable struct Scalar{T} - ptr::Ptr{Cvoid} - device::Int + ptr::Ptr{Cvoid} + device::Int end -function Scalar(r::T; dev = -1) where T <: Real - ptr = Ref(Ptr{Cvoid}()) +function Scalar(r::T; dev = -1) where {T <: Real} + ptr = Ref(Ptr{Cvoid}()) - ats_float(ptr, float.(r)) - Scalar{T}(ptr[], dev) + ats_float(ptr, float.(r)) + return Scalar{T}(ptr[], dev) end free!(s::Scalar) = ats_free(s.ptr) diff --git a/src/statistics.jl b/src/statistics.jl index 1e8e217..2cf2a54 100644 --- a/src/statistics.jl +++ b/src/statistics.jl @@ -1,30 +1,30 @@ using Statistics -function Statistics.mean(t::Tensor{T,N}; dims = :) where {T,N} - ptr = Ref(Ptr{Cvoid}()) +function Statistics.mean(t::Tensor{T, N}; dims = :) where {T, N} + ptr = Ref(Ptr{Cvoid}()) - if dims isa Colon - atg_mean(ptr, t.ptr, options[T]) - Tensor{T,0}(ptr[], on(t)) - else - atg_mean1(ptr, t.ptr, dims, length(dims), dims[1], options[T]) - Tensor{T,N-length(dims)}(ptr[], on(t)) - end + if dims isa Colon + atg_mean(ptr, t.ptr, options[T]) + Tensor{T, 0}(ptr[], on(t)) + else + atg_mean1(ptr, t.ptr, dims, length(dims), dims[1], options[T]) + Tensor{T, N - length(dims)}(ptr[], on(t)) + end end -function _sum(t::Tensor{T}, ::Colon) where T - ptr = Ref(Ptr{Cvoid}()) +function _sum(t::Tensor{T}, ::Colon) where {T} + ptr = Ref(Ptr{Cvoid}()) - atg_sum(ptr, t.ptr, options[T]) - Tensor{T,0}(ptr[], on(t)) |> collect |> x -> x[] + atg_sum(ptr, t.ptr, options[T]) + return Tensor{T, 0}(ptr[], on(t)) |> collect |> x -> x[] end -function _sum(t::Tensor{T,N}, dims) where {T,N} - ptr = Ref(Ptr{Cvoid}()) +function _sum(t::Tensor{T, N}, dims) where {T, N} + ptr = Ref(Ptr{Cvoid}()) - dims = N .- dims - atg_sum1(ptr, t.ptr, [dims], length(dims), 1, options[T]) - Tensor{T,N}(ptr[], on(t)) + dims = N .- dims + atg_sum1(ptr, t.ptr, [dims], length(dims), 1, options[T]) + return Tensor{T, N}(ptr[], on(t)) end Statistics.sum(t::Tensor; dims = :) = _sum(t, dims) diff --git a/src/tensor.jl b/src/tensor.jl index 71401f2..6e2e386 100644 --- a/src/tensor.jl +++ b/src/tensor.jl @@ -1,84 +1,81 @@ -const options = Dict( - Int32 => 3, - Int64 => 4, - Float32 => 6, - Float64 => 7,) +const options = Dict(Int32 => 3, Int64 => 4, Float32 => 6, Float64 => 7) let was_enabled = Ref{Int32}() - at_grad_set_enabled(was_enabled, 0) + at_grad_set_enabled(was_enabled, 0) end struct TorchGPUOOMError <: Exception end function no_grad(f; flag = 0) - at_no_grad(flag) - f() + at_no_grad(flag) + return f() end -async_free!(x) = let x = x, ptr = x.ptr, oid = objectid(x) - @async begin - free!(x) - end - return -end +async_free!(x) = + let x = x, ptr = x.ptr, oid = objectid(x) + @async begin + free!(x) + end + return nothing + end -mutable struct Tensor{T, N} <: AbstractArray{T,N} - ptr::Ptr - device::Int +mutable struct Tensor{T, N} <: AbstractArray{T, N} + ptr::Ptr + device::Int - function Tensor{T,N}(ptr::Ptr, dev::Int) where {T,N} - obj = new(ptr, dev) - finalizer(async_free!, obj) - # TURN_ON_LOGGING == true && (logdict[ptr] = (size(obj), stacktrace())) - obj - end + function Tensor{T, N}(ptr::Ptr, dev::Int) where {T, N} + obj = new(ptr, dev) + finalizer(async_free!, obj) + # TURN_ON_LOGGING == true && (logdict[ptr] = (size(obj), stacktrace())) + return obj + end end TensorVector{T} = Tensor{T, 1} TensorMatrix{T} = Tensor{T, 2} TensorVecOrMat{T} = Union{TensorVector{T}, TensorMatrix{T}} -function Tensor(::Type{T}, sz::Int...; dev = -1) where T - ptr = Ref(Ptr{Cvoid}()) - dtype = options[T] - sz = length(sz) == 2 ? collect(sz) : reverse(collect(sz)) - mem = dev - d = Ref(pointer(sz)) - len = length(sz) +function Tensor(::Type{T}, sz::Int...; dev = -1) where {T} + ptr = Ref(Ptr{Cvoid}()) + dtype = options[T] + sz = length(sz) == 2 ? collect(sz) : reverse(collect(sz)) + mem = dev + d = Ref(pointer(sz)) + len = length(sz) - # atg_rand - atg_zeros(ptr, d.x, len, dtype, mem) - Tensor{T, len}(ptr[], dev) + # atg_rand + atg_zeros(ptr, d.x, len, dtype, mem) + return Tensor{T, len}(ptr[], dev) end -Tensor(sz::Int...; dev = -1) = Tensor(Float32, sz..., dev = dev) -Tensor(sz::Int; dev = -1) = Tensor(Float32, Int(sz), dev = dev) +Tensor(sz::Int...; dev = -1) = Tensor(Float32, sz...; dev = dev) +Tensor(sz::Int; dev = -1) = Tensor(Float32, Int(sz); dev = dev) # function Tensor{T,N}(ptr::Ptr) where {T,N} # Tensor{T,N}(ptr, on(ptr)) # end function Base.ndims(t::Tensor) - i = Int32[-1] - at_dim(i, t.ptr) - Int(i[1]) + i = Int32[-1] + at_dim(i, t.ptr) + return Int(i[1]) end function Base.size(t::Tensor) - dims = ndims(t) - sz = zeros(Int32, dims) - at_shape(t.ptr, pointer(sz)) - # s = Int.(tuple(sz...)) - if t isa TensorMatrix - Int.(tuple(sz...)) - else - reverse(Int.(tuple(sz...))) - end + dims = ndims(t) + sz = zeros(Int32, dims) + at_shape(t.ptr, pointer(sz)) + # s = Int.(tuple(sz...)) + if t isa TensorMatrix + Int.(tuple(sz...)) + else + reverse(Int.(tuple(sz...))) + end end function Base.size(t::Tensor, dim::Int) - sz = size(t) - dim <= length(sz) ? sz[dim] : 1 + sz = size(t) + return dim <= length(sz) ? sz[dim] : 1 end Base.length(t::Tensor) = prod(size(t)) @@ -91,75 +88,72 @@ Base.IndexStyle(::Type{<:Tensor}) = IndexCartesian() # end function Base.similar(t::Tensor, ::Type{K}, sz::Int...) where {K} - Tensor(K, sz..., dev = on(t)) + return Tensor(K, sz...; dev = on(t)) end -Base.similar(t::Tensor{T,N}) where {T,N} = Tensor(T,size(t)..., dev = on(t)) -Base.similar(t::Tensor{T,N}, sz::Int...) where {T,N} = similar(t, T, sz...) +Base.similar(t::Tensor{T, N}) where {T, N} = Tensor(T, size(t)...; dev = on(t)) +Base.similar(t::Tensor{T, N}, sz::Int...) where {T, N} = similar(t, T, sz...) Base.similar(t::Tensor, dims::Tuple) = similar(t, dims...) -function Base.copy(t::Tensor{T,N}) where {T,N} - sz = size(t) - z = zeros(T, sz...) - copyto!(z, t) +function Base.copy(t::Tensor{T, N}) where {T, N} + sz = size(t) + z = zeros(T, sz...) + return copyto!(z, t) end function Base.copyto!(dest::AbstractArray, src::Tensor) - at_copy_data(src.ptr, dest, length(dest), sizeof(eltype(dest))) - dest + at_copy_data(src.ptr, dest, length(dest), sizeof(eltype(dest))) + return dest end Base.copyto!(dest::Tensor, src::Tensor) = at_copy_(dest.ptr, src.ptr) -function Base.reshape(t::Tensor{T,N}, dims::Union{Colon, Int}...) where {T,N} - ptr = Ref(Ptr{Cvoid}()) +function Base.reshape(t::Tensor{T, N}, dims::Union{Colon, Int}...) where {T, N} + ptr = Ref(Ptr{Cvoid}()) - dims = Colon() in dims ? Base._reshape_uncolon(t, dims) : dims - dims = length(dims) == 2 ? collect(dims) : reverse(collect(dims)) - atg_reshape(ptr, t.ptr, dims, length(dims)) - Tensor{T,length(dims)}(ptr[], on(t)) + dims = Colon() in dims ? Base._reshape_uncolon(t, dims) : dims + dims = length(dims) == 2 ? collect(dims) : reverse(collect(dims)) + atg_reshape(ptr, t.ptr, dims, length(dims)) + return Tensor{T, length(dims)}(ptr[], on(t)) end -function Base.zero(t::Tensor{T,N}) where {T, N} - ptr = Ref(Ptr{Cvoid}()) +function Base.zero(t::Tensor{T, N}) where {T, N} + ptr = Ref(Ptr{Cvoid}()) - atg_zeros_like(ptr, t.ptr) - Tensor{T, N}(ptr[], on(t)) + atg_zeros_like(ptr, t.ptr) + return Tensor{T, N}(ptr[], on(t)) end -function Base.rand(::Type{Tensor{T}}, sz::Int...; dev = -1) where T <: Real - - ptr = Ref(Ptr{Cvoid}()) - dtype = options[T] - sz = collect(sz) - mem = dev - len = length(sz) +function Base.rand(::Type{Tensor{T}}, sz::Int...; dev = -1) where {T <: Real} + ptr = Ref(Ptr{Cvoid}()) + dtype = options[T] + sz = collect(sz) + mem = dev + len = length(sz) - # atg_rand - atg_rand(ptr, sz, len, dtype, mem) - Tensor{T, len}(ptr[], dev) + # atg_rand + atg_rand(ptr, sz, len, dtype, mem) + return Tensor{T, len}(ptr[], dev) end -Base.zeros(::Type{Tensor{T}}, sz::Int...; dev = -1) where T = - Tensor(T, sz..., dev = dev) - -function tensor(x::AbstractArray{T,N}; dev = -1) where {T,N} +Base.zeros(::Type{Tensor{T}}, sz::Int...; dev = -1) where {T} = Tensor(T, sz...; dev = dev) - sz = if N == 2 - collect(size(x)) - elseif N == 1 - [collect(size(x));1] - else - collect(size(x)) |> reverse - end - # d = Ref(pointer(sz)) - el_sz_in_bytes = sizeof(eltype(x)) - nd = ndims(x) - typ = options[T] - parr = Ref(pointer(x)) +function tensor(x::AbstractArray{T, N}; dev = -1) where {T, N} + sz = if N == 2 + collect(size(x)) + elseif N == 1 + [collect(size(x)); 1] + else + collect(size(x)) |> reverse + end + # d = Ref(pointer(sz)) + el_sz_in_bytes = sizeof(eltype(x)) + nd = ndims(x) + typ = options[T] + parr = Ref(pointer(x)) - ptr = Ref(Ptr{Cvoid}()) - at_tensor_of_data(ptr, parr.x, sz, nd, el_sz_in_bytes, typ) - opt = Tensor{Float32, N}(ptr[], dev) - to(opt, dev = dev) + ptr = Ref(Ptr{Cvoid}()) + at_tensor_of_data(ptr, parr.x, sz, nd, el_sz_in_bytes, typ) + opt = Tensor{Float32, N}(ptr[], dev) + return to(opt; dev = dev) end # tensor(x) = x tensor(x::Fill; kwargs...) = tensor(collect(x); kwargs...) @@ -168,24 +162,24 @@ tensor(x::Tensor; kwargs...) = x Base.print_array(io::IO, t::Tensor) = Base.print_array(io, collect(t)) Base.show_vector(io::IO, t::Tensor) = Base.show_vector(io, collect(t)) -function from_blob(x::AbstractArray{T,N}; dev = -1) where {T,N} - sz = reverse(collect(size(x))) - st = reverse(collect(strides(x))) - ptr = Ref(Ptr{Cvoid}()) - at_from_blob(ptr, pointer(x), sz, length(sz), st, length(st), dev) - Tensor{T,N}(ptr[], dev) +function from_blob(x::AbstractArray{T, N}; dev = -1) where {T, N} + sz = reverse(collect(size(x))) + st = reverse(collect(strides(x))) + ptr = Ref(Ptr{Cvoid}()) + at_from_blob(ptr, pointer(x), sz, length(sz), st, length(st), dev) + return Tensor{T, N}(ptr[], dev) end -function to(x::Tensor{T,N}; dev = -1) where {T,N} - ptr = Ref(Ptr{Cvoid}()) - atg_to(ptr, x.ptr, dev) - Tensor{Float32,N}(ptr[], dev) +function to(x::Tensor{T, N}; dev = -1) where {T, N} + ptr = Ref(Ptr{Cvoid}()) + atg_to(ptr, x.ptr, dev) + return Tensor{Float32, N}(ptr[], dev) end on(t::Tensor) = t.device function free!(t::Tensor) - # TURN_ON_LOGGING && delete!(logdict, t.ptr) - at_free(t.ptr) + # TURN_ON_LOGGING && delete!(logdict, t.ptr) + return at_free(t.ptr) end free!(ptr::Ptr) = at_free(ptr) diff --git a/test/flux_tests.jl b/test/flux_tests.jl index 9923a56..84746e5 100644 --- a/test/flux_tests.jl +++ b/test/flux_tests.jl @@ -15,7 +15,7 @@ end tresnet = Flux.fmap(x -> Torch.to_tensor(x; dev = torch_device), resnet.layers) ip = rand(Float32, 224, 224, 3, 1) # An RGB Image - tip = tensor(ip, dev = torch_device) + tip = tensor(ip; dev = torch_device) top = tresnet(tip) op = resnet.layers(ip) diff --git a/test/runtests.jl b/test/runtests.jl index 8547b05..336ee19 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,6 +1,6 @@ using Test -@testset verbose=true "Torch" begin +@testset verbose = true "Torch" begin include("flux_tests.jl") include("tensor_movement_tests.jl") include("tensor_nnlib_tests.jl") diff --git a/test/tensor_movement_tests.jl b/test/tensor_movement_tests.jl index 041da9f..242e4d2 100644 --- a/test/tensor_movement_tests.jl +++ b/test/tensor_movement_tests.jl @@ -8,8 +8,8 @@ else end @testset "Movement" begin - r = rand(Float32, 3,3) - tr = tensor(r, dev = torch_device) + r = rand(Float32, 3, 3) + tr = tensor(r; dev = torch_device) @test tr isa Tensor @test tr .* tr isa Tensor diff --git a/test/tensor_nnlib_tests.jl b/test/tensor_nnlib_tests.jl index e68dabd..8690275 100644 --- a/test/tensor_nnlib_tests.jl +++ b/test/tensor_nnlib_tests.jl @@ -16,15 +16,22 @@ end kernel = rand(-9.0f0:9.0f0, kernel_height, kernel_width, 1, in_channels) - for height in [5, 6], - width in [5, 7] - + for height in [5, 6], width in [5, 7] test_input = rand(-9.0f0:9.0f0, height, width, in_channels, 1) - x = tensor(test_input, dev = torch_device) - w = tensor(kernel, dev = torch_device) - - expected_output = NNlib.depthwiseconv(test_input, kernel, pad = (0,0), stride = (1,1 ), dilation = (1, 1), flipped = true) - test_output = NNlib.depthwiseconv(x, w, pad = (0,0), stride = (1,1 ), dilation = (1, 1)) + x = tensor(test_input; dev = torch_device) + w = tensor(kernel; dev = torch_device) + + expected_output = NNlib.depthwiseconv( + test_input, + kernel; + pad = (0, 0), + stride = (1, 1), + dilation = (1, 1), + flipped = true, + ) + test_output = NNlib.depthwiseconv( + x, w; pad = (0, 0), stride = (1, 1), dilation = (1, 1) + ) test_output = Array(test_output) @test maximum(abs.(test_output - expected_output)) < 1e3 * eps(Float32) # Numerical accuracy adjusted wrt. CI @@ -32,7 +39,6 @@ end end end - @testset "Conv with padding" begin for kernel_width in [1, 2, 3, 5], kernel_height in [1, 2, 3, 5], @@ -40,27 +46,29 @@ end out_channels in [1, 2] num_coefficients = (kernel_width * kernel_height * in_channels * out_channels) - kernel = reshape(1.0f0:num_coefficients, kernel_height, kernel_width, in_channels, out_channels) + kernel = reshape( + 1.0f0:num_coefficients, kernel_height, kernel_width, in_channels, out_channels + ) kernel = collect(kernel) pad = size(kernel)[1:2] .÷ 2 - for height in [1, 2, 3, 4], - width in [1, 2, 3, 5] - + for height in [1, 2, 3, 4], width in [1, 2, 3, 5] test_input = zeros(Float32, height, width, in_channels, 1) test_input[(height + 1) ÷ 2, (width + 1) ÷ 2, 1, 1] = 1 - x = tensor(test_input, dev = torch_device) - w = tensor(kernel, dev = torch_device) - - cdims = NNlib.DenseConvDims(size(test_input), - size(kernel), - stride=(1, 1), - padding=pad, - dilation=(1, 1), - flipkernel = true) + x = tensor(test_input; dev = torch_device) + w = tensor(kernel; dev = torch_device) + + cdims = NNlib.DenseConvDims( + size(test_input), + size(kernel); + stride = (1, 1), + padding = pad, + dilation = (1, 1), + flipkernel = true, + ) expected_output = NNlib.conv(test_input, kernel, cdims) - test_output = NNlib.conv(x, w, cdims) + test_output = NNlib.conv(x, w, cdims) test_output = Array(test_output) @test maximum(abs.(test_output - expected_output)) < 1e6 * eps(Float32) # Numerical accuracy adjusted wrt. CI @@ -68,7 +76,6 @@ end end end - @testset "Conv with stride" begin for kernel_width in [1, 3, 4], kernel_height in [1, 2, 5], @@ -80,24 +87,24 @@ end kernel = fill(1.0f0, kernel_height, kernel_width, in_channels, out_channels) kernel = collect(kernel) - for height in 13:(13 + row_stride - 1), - width in 15:(15 + column_stride - 1) - + for height in 13:(13 + row_stride - 1), width in 15:(15 + column_stride - 1) sz_in = [height, width, in_channels, 1] test_input = reshape(1.0f0:prod(sz_in), height, width, in_channels, 1) test_input = collect(test_input) - x = tensor(test_input, dev = torch_device) - w = tensor(kernel, dev = torch_device) - - cdims = NNlib.DenseConvDims(size(test_input), - size(kernel), - stride=(row_stride, column_stride), - padding=(0, 0), - dilation=(1, 1), - flipkernel = true) + x = tensor(test_input; dev = torch_device) + w = tensor(kernel; dev = torch_device) + + cdims = NNlib.DenseConvDims( + size(test_input), + size(kernel); + stride = (row_stride, column_stride), + padding = (0, 0), + dilation = (1, 1), + flipkernel = true, + ) expected_output = NNlib.conv(test_input, kernel, cdims) - test_output = NNlib.conv(x, w, cdims) + test_output = NNlib.conv(x, w, cdims) test_output = Array(test_output) @test maximum(abs.(test_output - expected_output)) < 1e4 * eps(Float32) # Numerical accuracy adjusted wrt. CI @@ -105,7 +112,6 @@ end end end - @testset "Conv with dilation" begin for kernel_width in 1, kernel_height in 1:9, @@ -123,24 +129,24 @@ end kernel = fill(1.0f0, kernel_height, kernel_width, in_channels, out_channels) kernel = collect(kernel) - for height in 13:(13 + row_stride - 1), - width in [1] - + for height in 13:(13 + row_stride - 1), width in [1] sz_in = [height, width, in_channels, 1] test_input = reshape(1.0f0:prod(sz_in), height, width, in_channels, 1) test_input = collect(test_input) - x = tensor(test_input, dev = torch_device) - w = tensor(kernel, dev = torch_device) - - cdims = NNlib.DenseConvDims(size(test_input), - size(kernel), - stride=(row_stride, column_stride), - padding=(0, 0), - dilation=(1, 1), - flipkernel = true) + x = tensor(test_input; dev = torch_device) + w = tensor(kernel; dev = torch_device) + + cdims = NNlib.DenseConvDims( + size(test_input), + size(kernel); + stride = (row_stride, column_stride), + padding = (0, 0), + dilation = (1, 1), + flipkernel = true, + ) expected_output = NNlib.conv(test_input, kernel, cdims) - test_output = NNlib.conv(x, w, cdims) + test_output = NNlib.conv(x, w, cdims) test_output = Array(test_output) @test maximum(abs.(test_output - expected_output)) < 1e2 * eps(Float32) # Numerical accuracy adjusted wrt. CI @@ -148,7 +154,6 @@ end end end - @testset "Pooling" begin for fn in (NNlib.maxpool, NNlib.meanpool), column_span in 1:3, @@ -168,15 +173,17 @@ end channels in 1:2 test_input = rand(0.0f0:9.0f0, height, width, channels, 1) - x = tensor(test_input, dev = torch_device) + x = tensor(test_input; dev = torch_device) - pdims = NNlib.PoolDims(size(test_input), - (row_span, column_span), - padding=padding, - stride=(row_stride, column_stride)) + pdims = NNlib.PoolDims( + size(test_input), + (row_span, column_span); + padding = padding, + stride = (row_stride, column_stride), + ) expected_output = fn(test_input, pdims) - test_output = fn(x, pdims) + test_output = fn(x, pdims) test_output = Array(test_output) @test maximum(abs.(test_output - expected_output)) < 1e1 * eps(Float32) @@ -184,7 +191,6 @@ end end end - @testset "Activations" begin for fn in (NNlib.relu, NNlib.tanh, NNlib.sigmoid, NNlib.leakyrelu, NNlib.softmax), height in [1, 2, 3, 4, 7], @@ -192,14 +198,14 @@ end channels in 1:3 test_input = rand(-9.0f0:9.0f0, height, width, channels, 1) - x = tensor(test_input, dev = torch_device) + x = tensor(test_input; dev = torch_device) if fn == NNlib.softmax - expected_output = fn(test_input, dims = 3) - test_output = fn(x, dims = 3) + expected_output = fn(test_input; dims = 3) + test_output = fn(x; dims = 3) else expected_output = fn.(test_input) - test_output = fn(x) + test_output = fn(x) end test_output = Array(test_output)