DNN/CNN/ContextHandles.h

0001 // @(#)root/tmva/tmva/dnn:$Id$
0002 // Author: Simon Pfreundschuh 20/06/16
0003
0004 /*************************************************************************
0005  * Copyright (C) 2016, Simon Pfreundschuh                                *
0006  * All rights reserved.                                                  *
0007  *                                                                       *
0008  * For the licensing terms see $ROOTSYS/LICENSE.                         *
0009  * For the list of contributors see $ROOTSYS/README/CREDITS.             *
0010  *************************************************************************/
0011
0012 /////////////////////////////////////////////////////////////////////
0013 // Contains function enums for activation and output functions, as //
0014 // well as generic evaluation functions, that delegate the call to //
0015 // the corresponding evaluation kernel.                            //
0016 /////////////////////////////////////////////////////////////////////
0017
0018 #ifndef TMVA_DNN_CNN_DESCRIPTORS
0019 #define TMVA_DNN_CNN_DESCRIPTORS
0020
0021 #include <stddef.h>
0022 #include <vector>
0023
0024 namespace TMVA
0025 {
0026 namespace DNN
0027 {
0028
0029 struct TDescriptors {
0030    virtual ~TDescriptors() {}
0031 };
0032 struct TWorkspace {
0033     virtual ~TWorkspace() {}
0034 };
0035
0036 template <typename Layer_t>
0037 struct TDNNGenDescriptors : public TMVA::DNN::TDescriptors {
0038    using HelperDescriptor_t = typename Layer_t::HelperDescriptor_t;
0039
0040    HelperDescriptor_t HelperDescriptor;
0041 };
0042
0043 namespace CNN {
0044
0045 //______________________________________________________________________________
0046 //
0047 // Keeps the descriptors for the CNN
0048 //______________________________________________________________________________
0049
0050 template <typename Layer_t>
0051 struct TCNNDescriptors : public TMVA::DNN::TDescriptors {
0052    using LayerDescriptor_t = typename Layer_t::LayerDescriptor_t;     ///< Main layer operation
0053    using HelperDescriptor_t = typename Layer_t::HelperDescriptor_t;   ///< Used to define possible helpers for the layers (e.g. activations)
0054    using WeightsDescriptor_t = typename Layer_t::WeightsDescriptor_t; ///< The weights that are modified (e.g filters)
0055
0056    LayerDescriptor_t LayerDescriptor;
0057    HelperDescriptor_t HelperDescriptor;
0058    WeightsDescriptor_t WeightsDescriptor;
0059 };
0060
0061 template <typename Layer_t>
0062 struct TCNNWorkspace : public TMVA::DNN::TWorkspace {
0063    using AlgorithmForward_t = typename Layer_t::AlgorithmForward_t;   ///< Forward layer operation
0064    using AlgorithmBackward_t = typename Layer_t::AlgorithmBackward_t; ///< Backward layer operation
0065    using AlgorithmHelper_t = typename Layer_t::AlgorithmHelper_t;     ///< Used for weight grad backward pass
0066
0067    using ReduceTensorDescriptor_t = typename Layer_t::ReduceTensorDescriptor_t;
0068
0069    using AlgorithmDataType_t = typename Layer_t::AlgorithmDataType_t;
0070
0071    AlgorithmForward_t AlgorithmForward;
0072    AlgorithmBackward_t AlgorithmBackward;
0073    AlgorithmHelper_t HelperAlgorithm;
0074
0075    AlgorithmDataType_t DataType;
0076
0077    size_t *ForwardWorkspace;
0078    size_t *BackwardWorkspace;
0079    size_t *HelperWorkspace;
0080
0081    void *fReductionWorkspace = nullptr;
0082
0083    size_t ForwardWorkspaceSize;
0084    size_t BackwardWorkspaceSize;
0085    size_t HelperWorkspaceSize;
0086    size_t fReductionWorkspaceSize = 0;
0087
0088    ReduceTensorDescriptor_t fReduceTensorDesc;
0089 };
0090
0091 } // namespace CNN
0092
0093 namespace RNN {
0094 template <typename Architecture_t>
0095 struct TRNNDescriptors : public TMVA::DNN::TDescriptors {
0096
0097    using LayerDescriptor_t = typename Architecture_t::RecurrentDescriptor_t; ///< Main layer operation
0098    using TensorDescriptor_t = typename Architecture_t::TensorDescriptor_t;   ///< the vector of tensor descriptors
0099    using HelperDescriptor_t = typename Architecture_t::DropoutDescriptor_t;  ///< use for dropout
0100
0101    LayerDescriptor_t LayerDescriptor;
0102    HelperDescriptor_t HelperDescriptor;
0103
0104 #if (CUDNN_VERSION >= 8000)
0105    //using WeightsDescriptor_t = typename Architecture_t::TensorDescriptor_t;
0106    using DataDescriptor_t = typename Architecture_t::RNNDataDescriptor_t;   ///< the vector of tensor descriptors
0107
0108    DataDescriptor_t xDataDesc;   // input data descriptor
0109    DataDescriptor_t yDataDesc;   // output data descriptor
0110
0111 #else
0112    using WeightsDescriptor_t = typename Architecture_t::FilterDescriptor_t;  ///< The weights that are modified (e.g filters)
0113    using DataDescriptor_t = typename Architecture_t::TensorDescriptor_t;   ///< the vector of tensor descriptors
0114
0115
0116    WeightsDescriptor_t WeightsDescriptor;
0117    WeightsDescriptor_t WeightsGradDescriptor;
0118
0119
0120    // for RNN need 4 vectors of tensor descriptors
0121
0122    std::vector<TensorDescriptor_t> xDesc;
0123    std::vector<TensorDescriptor_t> yDesc;
0124    std::vector<TensorDescriptor_t> dxDesc;
0125    std::vector<TensorDescriptor_t> dyDesc;
0126
0127 #endif
0128
0129 };
0130
0131 template <typename Layer_t>
0132 struct TRNNWorkspace : public TMVA::DNN::TWorkspace {
0133
0134    void *ForwardWorkspace = nullptr;
0135    void *InferenceWorkspace = nullptr;
0136    void *HelperWorkspace = nullptr;
0137
0138
0139    size_t ForwardWorkspaceSize = 0;
0140    size_t InferenceWorkspaceSize = 0;
0141    size_t HelperWorkspaceSize = 0;
0142 };
0143
0144 }  // end namespace RNN
0145
0146
0147 } // namespace DNN
0148 } // namespace TMVA
0149
0150 #endif