PKU-DAIR
/
Hetu
Not watched
1
0
Fork 0
Code
Releases 0 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Branch: master
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'master'
${ noResults }
Hetu/python/hetu/gpu_links/CudnnAvgPoolLink.py

22 lines
1.0 kB
Raw Permalink Blame History 

  1. from __future__ import absolute_import

  2. 

  3. import ctypes

  4. from .._base import _LIB

  5. from .. import ndarray as _nd

  6. 

  7. 

  8. def CuDNN_average_pooling2d(in_arr, kernel_H, kernel_W, pooled_layer, padding=0, stride=1, stream=None):

  9.     assert isinstance(in_arr, _nd.NDArray)

  10.     assert isinstance(pooled_layer, _nd.NDArray)

  11.     _LIB.CuDNN_DLGpuAvgerage_Pooling2d(

  12.         in_arr.handle, kernel_H, kernel_W, pooled_layer.handle, padding, stride, stream.handle if stream else None)

  13. 

  14. 

  15. def CuDNN_average_pooling2d_gradient(in_arr_y, in_gradient_y, in_arr_x, kernel_H, kernel_W, out_gradient_x, padding=0, stride=1, stream=None):

  16.     assert isinstance(in_arr_y, _nd.NDArray)

  17.     assert isinstance(in_gradient_y, _nd.NDArray)

  18.     assert isinstance(in_arr_x, _nd.NDArray)

  19.     assert isinstance(out_gradient_x, _nd.NDArray)

  20.     _LIB.CuDNN_DLGpuAvgerage_Pooling2d_gradient(

  21.         in_arr_y.handle, in_gradient_y.handle, in_arr_x.handle, kernel_H, kernel_W, out_gradient_x.handle, padding, stride, stream.handle if stream else None)