argmax layer ND-tensor

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

src/layers/argmax.jl

@@ -1,45 +1,52 @@
 @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)),
 )
 
 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

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

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)