<s> import atexit <EOL> import numpy <EOL> import six <EOL> from cupy import cuda <EOL> from cupy . cuda import cudnn <EOL> _cudnn_version = cudnn . getVersion ( ) <EOL> _handles = { } <EOL> def get_handle ( ) : <EOL> global _handles <EOL> device = cuda . Device ( ) <EOL> handle = _handles . get ( device . id , None ) <EOL> if handle is None : <EOL> handle = cudnn . create ( ) <EOL> _handles [ device . id ] = handle <EOL> return handle <EOL> @ atexit . register <EOL> def reset_handles ( ) : <EOL> global _handles <EOL> handles = _handles <EOL> _handles = { } <EOL> for handle in six . itervalues ( handles ) : <EOL> cudnn . destroy ( handle ) <EOL> class Descriptor ( object ) : <EOL> def __init__ ( self , descriptor , destroyer ) : <EOL> self . value = descriptor <EOL> self . destroy = destroyer <EOL> def __del__ ( self ) : <EOL> if self . value : <EOL> self . destroy ( self . value ) <EOL> self . value = None <EOL> def get_data_type ( dtype ) : <EOL> if dtype . type == numpy . float32 : <EOL> return cudnn . CUDNN_DATA_FLOAT <EOL> elif dtype . type == numpy . float64 : <EOL> return cudnn . CUDNN_DATA_DOUBLE <EOL> elif dtype . type == numpy . float16 : <EOL> return cudnn . CUDNN_DATA_HALF <EOL> else : <EOL> raise TypeError ( '<STR_LIT>' . format ( dtype ) ) <EOL> def _to_ctypes_array ( tup , dtype = numpy . intc ) : <EOL> return numpy . array ( tup , dtype = dtype ) . ctypes <EOL> def create_tensor_descriptor ( arr , format = cudnn . CUDNN_TENSOR_NCHW ) : <EOL> desc = Descriptor ( cudnn . createTensorDescriptor ( ) , <EOL> cudnn . destroyTensorDescriptor ) <EOL> if arr . ndim != <NUM_LIT:4> : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> if not arr . flags . c_contiguous : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> data_type = get_data_type ( arr . dtype ) <EOL> cudnn . setTensor4dDescriptor ( desc . value , format , data_type , <EOL> * arr . shape ) <EOL> return desc <EOL> def create_filter_descriptor ( arr , mode = cudnn . CUDNN_CROSS_CORRELATION ) : <EOL> desc = Descriptor ( cudnn . createFilterDescriptor ( ) , <EOL> cudnn . destroyFilterDescriptor ) <EOL> data_type = get_data_type ( arr . dtype ) <EOL> if arr . ndim == <NUM_LIT:4> : <EOL> cudnn . setFilter4dDescriptor_v3 ( desc . value , data_type , * arr . shape ) <EOL> else : <EOL> c_shape = _to_ctypes_array ( arr . shape ) <EOL> cudnn . setFilterNdDescriptor_v3 ( desc . value , data_type , arr . ndim , <EOL> c_shape . data ) <EOL> return desc <EOL> def create_convolution_descriptor ( pad , stride , <EOL> mode = cudnn . CUDNN_CROSS_CORRELATION ) : <EOL> desc = Descriptor ( cudnn . createConvolutionDescriptor ( ) , <EOL> cudnn . destroyConvolutionDescriptor ) <EOL> ndim = len ( pad ) <EOL> if ndim != len ( stride ) : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> if ndim == <NUM_LIT:2> : <EOL> cudnn . setConvolution2dDescriptor ( <EOL> desc . value , pad [ <NUM_LIT:0> ] , pad [ <NUM_LIT:1> ] , stride [ <NUM_LIT:0> ] , stride [ <NUM_LIT:1> ] , <NUM_LIT:1> , <NUM_LIT:1> , mode ) <EOL> else : <EOL> c_pad = _to_ctypes_array ( pad ) <EOL> c_stride = _to_ctypes_array ( stride ) <EOL> c_upscale = _to_ctypes_array ( ( <NUM_LIT:1> , ) * ndim ) <EOL> cudnn . setConvolutionNdDescriptor_v2 ( <EOL> desc . value , ndim , c_pad . data , c_stride . data , c_upscale . data , mode ) <EOL> return desc <EOL> def create_pooling_descriptor ( ksize , stride , pad , mode ) : <EOL> desc = Descriptor ( cudnn . createPoolingDescriptor ( ) , <EOL> cudnn . destroyPoolingDescriptor ) <EOL> ndim = len ( ksize ) <EOL> if ndim != len ( stride ) or ndim != len ( pad ) : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> if ndim == <NUM_LIT:2> : <EOL> cudnn . setPooling2dDescriptor_v3 ( <EOL> desc . value , mode , ksize [ <NUM_LIT:0> ] , ksize [ <NUM_LIT:1> ] , pad [ <NUM_LIT:0> ] , pad [ <NUM_LIT:1> ] , <EOL> stride [ <NUM_LIT:0> ] , stride [ <NUM_LIT:1> ] ) <EOL> else : <EOL> c_ksize = _to_ctypes_array ( ksize ) <EOL> c_pad = _to_ctypes_array ( pad ) <EOL> c_stride = _to_ctypes_array ( stride ) <EOL> cudnn . setPoolingNdDescriptor_v3 ( <EOL> desc . value , mode , ndim , c_ksize . data , c_pad . data , c_stride . data ) <EOL> return desc <EOL> if _cudnn_version >= <NUM_LIT> : <EOL> def add_tensor ( handle , alpha , biasDesc , biasData , beta , srcDestDesc , <EOL> srcDestData ) : <EOL> cudnn . addTensor_v3 ( handle , alpha , biasDesc , <EOL> biasData , beta , srcDestDesc , srcDestData ) <EOL> else : <EOL> def add_tensor ( handle , alpha , biasDesc , biasData , beta , srcDestDesc , <EOL> srcDestData ) : <EOL> cudnn . addTensor_v2 ( handle , cudnn . CUDNN_ADD_SAME_C , alpha , biasDesc , <EOL> biasData , beta , srcDestDesc , srcDestData ) </s>
