|
|
#ifdef USE_CUDNN |
|
|
#include <vector> |
|
|
|
|
|
#include "caffe/layers/cudnn_sigmoid_layer.hpp" |
|
|
|
|
|
namespace caffe { |
|
|
|
|
|
template <typename Dtype> |
|
|
void CuDNNSigmoidLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom, |
|
|
const vector<Blob<Dtype>*>& top) { |
|
|
SigmoidLayer<Dtype>::LayerSetUp(bottom, top); |
|
|
|
|
|
CUDNN_CHECK(cudnnCreate(&handle_)); |
|
|
cudnn::createTensor4dDesc<Dtype>(&bottom_desc_); |
|
|
cudnn::createTensor4dDesc<Dtype>(&top_desc_); |
|
|
cudnn::createActivationDescriptor<Dtype>(&activ_desc_, |
|
|
CUDNN_ACTIVATION_SIGMOID); |
|
|
handles_setup_ = true; |
|
|
} |
|
|
|
|
|
template <typename Dtype> |
|
|
void CuDNNSigmoidLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom, |
|
|
const vector<Blob<Dtype>*>& top) { |
|
|
SigmoidLayer<Dtype>::Reshape(bottom, top); |
|
|
const int N = bottom[0]->num(); |
|
|
const int K = bottom[0]->channels(); |
|
|
const int H = bottom[0]->height(); |
|
|
const int W = bottom[0]->width(); |
|
|
cudnn::setTensor4dDesc<Dtype>(&bottom_desc_, N, K, H, W); |
|
|
cudnn::setTensor4dDesc<Dtype>(&top_desc_, N, K, H, W); |
|
|
} |
|
|
|
|
|
template <typename Dtype> |
|
|
CuDNNSigmoidLayer<Dtype>::~CuDNNSigmoidLayer() { |
|
|
|
|
|
if (!handles_setup_) { return; } |
|
|
|
|
|
cudnnDestroyTensorDescriptor(this->bottom_desc_); |
|
|
cudnnDestroyTensorDescriptor(this->top_desc_); |
|
|
cudnnDestroy(this->handle_); |
|
|
} |
|
|
|
|
|
INSTANTIATE_CLASS(CuDNNSigmoidLayer); |
|
|
|
|
|
} |
|
|
#endif |
|
|
|