<s> import numpy <EOL> from chainer import cuda <EOL> from chainer import function <EOL> from chainer . utils import type_check <EOL> if cuda . cudnn_enabled : <EOL> cudnn = cuda . cudnn <EOL> libcudnn = cudnn . cudnn <EOL> _algorithm = libcudnn . CUDNN_SOFTMAX_ACCURATE <EOL> _mode = libcudnn . CUDNN_SOFTMAX_MODE_CHANNEL <EOL> class Softmax ( function . Function ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , use_cudnn = True ) : <EOL> self . use_cudnn = use_cudnn <EOL> def check_type_forward ( self , in_types ) : <EOL> type_check . expect ( in_types . size ( ) == <NUM_LIT:1> ) <EOL> x_type , = in_types <EOL> type_check . expect ( <EOL> x_type . dtype == numpy . float32 , <EOL> x_type . ndim > <NUM_LIT:1> , <EOL> ) <EOL> def forward ( self , x ) : <EOL> xp = cuda . get_array_module ( * x ) <EOL> if xp != numpy and cuda . cudnn_enabled and self . use_cudnn : <EOL> dtype = x [ <NUM_LIT:0> ] . dtype <EOL> one = numpy . array ( <NUM_LIT:1> , dtype = dtype ) . ctypes <EOL> zero = numpy . array ( <NUM_LIT:0> , dtype = dtype ) . ctypes <EOL> handle = cudnn . get_handle ( ) <EOL> x_cube = x [ <NUM_LIT:0> ] . reshape ( x [ <NUM_LIT:0> ] . shape [ : <NUM_LIT:2> ] + ( - <NUM_LIT:1> , <NUM_LIT:1> ) ) <EOL> desc = cudnn . create_tensor_descriptor ( x_cube ) <EOL> self . y = xp . empty_like ( x [ <NUM_LIT:0> ] ) <EOL> libcudnn . softmaxForward ( <EOL> handle , _algorithm , _mode , one . data , desc . value , <EOL> x_cube . data . ptr , zero . data , desc . value , <EOL> self . y . data . ptr ) <EOL> else : <EOL> self . y = x [ <NUM_LIT:0> ] - x [ <NUM_LIT:0> ] . max ( axis = <NUM_LIT:1> , keepdims = True ) <EOL> xp . exp ( self . y , out = self . y ) <EOL> self . y /= self . y . sum ( axis = <NUM_LIT:1> , keepdims = True ) <EOL> return self . y , <EOL> def backward ( self , x , gy ) : <EOL> xp = cuda . get_array_module ( * x ) <EOL> if xp != numpy and cuda . cudnn_enabled and self . use_cudnn : <EOL> dtype = x [ <NUM_LIT:0> ] . dtype <EOL> one = numpy . array ( <NUM_LIT:1> , dtype = dtype ) . ctypes <EOL> zero = numpy . array ( <NUM_LIT:0> , dtype = dtype ) . ctypes <EOL> handle = cudnn . get_handle ( ) <EOL> gx = xp . empty_like ( x [ <NUM_LIT:0> ] ) <EOL> gx_cube = gx . reshape ( gx . shape [ : <NUM_LIT:2> ] + ( - <NUM_LIT:1> , <NUM_LIT:1> ) ) <EOL> desc = cudnn . create_tensor_descriptor ( gx_cube ) <EOL> libcudnn . softmaxBackward ( <EOL> handle , _algorithm , _mode , one . data , desc . value , <EOL> self . y . data . ptr , desc . value , gy [ <NUM_LIT:0> ] . data . ptr , zero . data , <EOL> desc . value , gx . data . ptr ) <EOL> else : <EOL> gx = self . y * gy [ <NUM_LIT:0> ] <EOL> sumdx = gx . sum ( axis = <NUM_LIT:1> , keepdims = True ) <EOL> gx -= self . y * sumdx <EOL> return gx , <EOL> def softmax ( x , use_cudnn = True ) : <EOL> """<STR_LIT>""" <EOL> return Softmax ( use_cudnn ) ( x ) </s>
