-
Notifications
You must be signed in to change notification settings - Fork 254
/
softmax.jl
32 lines (29 loc) · 1.32 KB
/
softmax.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
type CuDNNSoftmaxState
inputs_desc :: Vector{CuDNN.Tensor4dDescriptor}
outputs_desc :: Vector{CuDNN.Tensor4dDescriptor}
end
function setup_etc(backend::GPUBackend, layer::SoftmaxLayer, dims::Vector{Int}, data_type, inputs)
inputs_desc = Array(CuDNN.Tensor4dDescriptor, length(inputs))
outputs_desc = Array(CuDNN.Tensor4dDescriptor, length(inputs))
for i = 1:length(inputs)
dim_sp, dim_prob, dim_num = split_dims(inputs[i], dims[i])
inputs_desc[i] = CuDNN.create_tensor4d_descriptor(data_type, (1,dim_sp,dim_prob,dim_num))
outputs_desc[i] = CuDNN.create_tensor4d_descriptor(data_type, (1,dim_sp,dim_prob,dim_num))
end
etc = CuDNNSoftmaxState(inputs_desc, outputs_desc)
return etc
end
function shutdown(backend::GPUBackend, state::SoftmaxLayerState)
map(destroy, state.blobs)
map(CuDNN.destroy_tensor4d_descriptor, state.etc.inputs_desc)
map(CuDNN.destroy_tensor4d_descriptor, state.etc.outputs_desc)
end
function forward(backend::GPUBackend, state::SoftmaxLayerState, inputs::Vector{Blob})
alpha = one(eltype(inputs[1]))
beta = zero(eltype(inputs[1]))
for i = 1:length(inputs)
CuDNN.softmax_forward(backend.cudnn_ctx, CuDNN.CUDNN_SOFTMAX_ACCURATE,
CuDNN.CUDNN_SOFTMAX_MODE_CHANNEL, alpha, state.etc.inputs_desc[i], inputs[i].ptr,
beta, state.etc.outputs_desc[i], state.blobs[i].ptr)
end
end