From 5fa3a4ecff203c9ceff8c663552f3d05bf8d853a Mon Sep 17 00:00:00 2001 From: Chiyuan Zhang Date: Fri, 19 Dec 2014 22:10:08 +0800 Subject: [PATCH] argmax layer ND-tensor --- src/cuda/layers/argmax.jl | 3 +-- src/layers/argmax.jl | 47 ++++++++++++++++++++++----------------- src/utils/tensor.jl | 6 ++--- test/layers/argmax.jl | 35 +++++++++++++++++------------ 4 files changed, 52 insertions(+), 39 deletions(-) diff --git a/src/cuda/layers/argmax.jl b/src/cuda/layers/argmax.jl index ae78eee..c2c435d 100644 --- a/src/cuda/layers/argmax.jl +++ b/src/cuda/layers/argmax.jl @@ -3,8 +3,7 @@ function forward(backend::GPUBackend, state::ArgmaxLayerState, inputs::Vector{Bl input = inputs[i] output = state.blobs[i] - width, height, channels, num = get_whcn(input) - spatial_dim = width*height + spatial_dim, channels, num = split_dims(input, state.dims[i]) data_type = eltype(input) x_block = int(ceil(float64(num)/CUDA.THREADS_PER_BLOCK_X)); diff --git a/src/layers/argmax.jl b/src/layers/argmax.jl index 481f0b6..a2b2ee8 100644 --- a/src/layers/argmax.jl +++ b/src/layers/argmax.jl @@ -1,5 +1,6 @@ @defstruct ArgmaxLayer Layer ( name :: String = "argmax", + (dim :: Int = -2, dim != 0), (tops :: Vector{Symbol} = Symbol[], length(tops) > 0), (bottoms :: Vector{Symbol} = Symbol[], length(bottoms) == length(tops)), ) @@ -7,39 +8,45 @@ type ArgmaxLayerState <: LayerState layer :: ArgmaxLayer blobs :: Vector{Blob} + + dims :: Vector{Int} end function setup(backend::Backend, layer::ArgmaxLayer, inputs::Vector{Blob}, diffs::Vector{Blob}) - blobs = map(inputs) do input - width, height, channels, num = get_whcn(input) - data_type = eltype(input) - - blob = make_blob(backend, data_type, width, height, 1, num) - blob + dims = Array(Int, length(inputs)) + blobs = Array(Blob, length(inputs)) + for i = 1:length(inputs) + total_dim = ndims(inputs[i]) + dim = layer.dim < 0 ? layer.dim + total_dim + 1 : layer.dim + @assert 1 <= dim <= total_dim + @assert dim != total_dim + dims[i] = dim + shape = [size(inputs[i])...] + shape[dim] = 1 + blobs[i] = make_blob(backend, eltype(inputs[i]), shape...) end - return ArgmaxLayerState(layer, blobs) + return ArgmaxLayerState(layer, blobs, dims) end function forward(backend::CPUBackend, state::ArgmaxLayerState, inputs::Vector{Blob}) for i = 1:length(inputs) input = inputs[i].data output = state.blobs[i].data - width, height, channels, num = get_whcn(input) - canonical_input = reshape(input, (width,height,channels,num)) - for n = 1:num - for w = 1:width - for h = 1:height - maxc = 1; maxval = canonical_input[w,h,maxc,n] - for c = 2:channels - @inbounds val = canonical_input[w,h,c,n] - if val > maxval - maxval = val - maxc = c - end + pre_dim, mid_dim, post_dim = split_dims(input, state.dims[i]) + for x = 0:pre_dim-1 + for z = 0:post_dim-1 + idx = Int[x + pre_dim*(y + mid_dim*z) for y=0:mid_dim-1] + 1 + maxc = 1 + @inbounds maxval = input[idx[1]] + for y = 2:length(idx) + @inbounds val = input[idx[y]] + if val > maxval + maxval = val + maxc = y end - @inbounds output[w,h,1,n] = maxc-1 end + @inbounds output[x + pre_dim*z + 1] = maxc-1 end end end diff --git a/src/utils/tensor.jl b/src/utils/tensor.jl index 0aa9e79..abbe1d6 100644 --- a/src/utils/tensor.jl +++ b/src/utils/tensor.jl @@ -4,9 +4,9 @@ export split_dims # (dim_pre, dim_mid, dim_post) function split_dims{T}(tensor::T, dim::Int) dims = size(tensor) - dim_pre = prod(dims[1:dim-1]) - dim_mid = dims[dim] - dim_post = prod(dims[dim+1:end]) + dim_pre ::Int = prod(dims[1:dim-1]) + dim_mid ::Int = dims[dim] + dim_post ::Int = prod(dims[dim+1:end]) (dim_pre, dim_mid, dim_post) end diff --git a/test/layers/argmax.jl b/test/layers/argmax.jl index 4ca3ca8..20acc4a 100644 --- a/test/layers/argmax.jl +++ b/test/layers/argmax.jl @@ -1,30 +1,32 @@ -function test_argmax_layer(backend::Backend, n_input, T, eps) +function test_argmax_layer(backend::Backend, n_input, tensor_dim, T, eps) println("-- Testing ArgmaxLayer on $(typeof(backend)){$T}...") - tensor_dim = abs(rand(Int)) % 4 + 2 println(" > $tensor_dim-dimensional tensor") dims = [abs(rand(Int, tensor_dim)) % 6 + 1 for i = 1:n_input] - input = [rand(T, dims[i]...) for i = 1:n_input] - input_blob = Blob[make_blob(backend, x) for x in input] + op_dim = max(abs(rand(Int)) % tensor_dim, 1) + inputs = [rand(T, dims[i]...) for i = 1:n_input] + input_blob = Blob[make_blob(backend, x) for x in inputs] diff_blob = Blob[NullBlob() for i = 1:n_input] println(" > Setup") - layer = ArgmaxLayer(bottoms=Array(Symbol,n_input),tops=Array(Symbol,n_input)) + layer = ArgmaxLayer(bottoms=Array(Symbol,n_input),tops=Array(Symbol,n_input),dim=op_dim) state = setup(backend, layer, input_blob, diff_blob) println(" > Forward") forward(backend, state, input_blob) for i = 1:n_input - width,height,channels,num = get_whcn(input[i]) - got_output = zeros(T, width, height, 1, num) - canonical_input = reshape(input[i], (width,height,channels,num)) - expected_output = similar(got_output) - for n = 1:num - for w = 1:width - for h = 1:height - expected_output[w,h,1,n] = indmax(canonical_input[w,h,:,n])-1 - end + outdim = [size(inputs[i])...] + outdim[op_dim] = 1 + got_output = zeros(T, outdim...) + expected_output = zeros(T, outdim...) + + pre_dim, mid_dim, post_dim = split_dims(inputs[i], op_dim) + input = reshape(inputs[i], pre_dim, mid_dim, post_dim) + output = reshape(expected_output, pre_dim, 1, post_dim) + for x = 1:pre_dim + for z = 1:post_dim + output[x,1,z] = indmax(input[x,:,z])-1 end end @@ -34,6 +36,11 @@ function test_argmax_layer(backend::Backend, n_input, T, eps) shutdown(backend, state) end +function test_argmax_layer(backend::Backend, n_input, T, eps) + for i = 2:6 + test_argmax_layer(backend, n_input, i, T, eps) + end +end function test_argmax_layer(backend::Backend) test_argmax_layer(backend, 3, Float64, 1e-10) test_argmax_layer(backend, 3, Float32, 1e-10)