def get_parameters(fn, handle, weight_buf): """Returns weight and bias tensors for each layer of the RNN. These tensors are views on the underlying weight buffer allocated by CuDNN. Note: for LSTM and GRU, which have multiple parameters of each type (4 and 3, respectively), these parameters are concatenated along the first dimension. These parameters are returned in a consistent order by CuDNN: (reset, forget, cell, outut) for LSTM (reset, input, new) for GRU Args: fn: The RNN function object holding the RNN state handle: a CuDNN handle weight_buf: a 1D tensor containing the CuDNN-allocated weight (or grad_weight) buffer Returns: parameters: [(weight_ih, weight_hh, bias_ih, bias_hh)*], with length equal to the num_layers. """ cudnn_methods = [ cudnn.lib.cudnnGetRNNLinLayerMatrixParams, cudnn.lib.cudnnGetRNNLinLayerBiasParams ] params = [] num_linear_layers = _num_linear_layers(fn) num_layers = fn.num_directions * fn.num_layers for layer in range(num_layers): layer_params = [] for cudnn_method in cudnn_methods: for linear_id in range(num_linear_layers): lin_layer_mat_desc = cudnn.FilterDescriptor() matrix_pointer = ctypes.c_void_p() check_error( cudnn_method(handle, fn.rnn_desc, layer, fn.x_descs[0], fn.w_desc, ctypes.c_void_p(weight_buf.data_ptr()), linear_id, lin_layer_mat_desc, ctypes.byref(matrix_pointer))) data_type = ctypes.c_int() format = ctypes.c_int() nb_dims = ctypes.c_int() min_dim = 3 filter_dim_a = torch.IntTensor(min_dim) check_error( cudnn.lib.cudnnGetFilterNdDescriptor( lin_layer_mat_desc, min_dim, ctypes.byref(data_type), ctypes.byref(format), ctypes.byref(nb_dims), ctypes.c_void_p(filter_dim_a.data_ptr()))) filter_dim_a.resize_(nb_dims.value) elem_size = cudnn._sizeofmap[fn.datatype] offset_bytes = (matrix_pointer.value - weight_buf.data_ptr()) assert (offset_bytes % elem_size == 0) offset = offset_bytes // elem_size # for all the RNN types provided by CUDNN, all the ih weights # are the same size and are allocated in a contiguous chunk # (same for the hh weights, and the ih and hh biases). # Since we're storing all the weights in a single tensor anyway, # might as well merge the CUDNN ones into a single tensor as well if linear_id == 0 or linear_id == num_linear_layers / 2: assert (filter_dim_a.prod() == filter_dim_a[0]) param = fn.weight_buf.new().set_( weight_buf.storage(), offset, filter_dim_a[0] * num_linear_layers // 2, filter_dim_a[2]) layer_params.append(param) else: assert (cur_offset == offset) cur_offset = offset + filter_dim_a[0] params.append(layer_params) return params
def init_weight_descriptor(fn, weight): w_desc = cudnn.FilterDescriptor() w_view = weight.view(-1, 1, 1) # seems that filters require >=3 dimensions w_desc.set(w_view) return w_desc