def init_dropout_descriptor(fn, handle):
dropout_desc = cudnn.DropoutDescriptor()
dropout_states_size = ctypes.c_long()
check_error(
cudnn.lib.cudnnDropoutGetStatesSize(handle,
ctypes.byref(dropout_states_size)))
dropout_states = torch.cuda.ByteTensor(dropout_states_size.value)
dropout_desc.set(handle, fn.dropout, dropout_states, fn.seed)
return dropout_desc
def init_dropout_descriptor(handle, dropout, train, dropout_seed,
dropout_state):
dropout_desc_name = 'desc_' + str(torch.cuda.current_device())
dropout_p = dropout if train else 0
if (dropout_desc_name not in dropout_state) or (
dropout_state[dropout_desc_name].get() is None):
dropout_state[dropout_desc_name] = Unserializable(
cudnn.DropoutDescriptor(handle, dropout_p, dropout_seed))
dropout_desc = dropout_state[dropout_desc_name].get()
dropout_desc.set_dropout(dropout_p, dropout_seed)
return dropout_desc
def init_rnn_descriptor(fn, handle):
dropout_desc_name = 'desc_' + str(torch.cuda.current_device())
dropout_p = fn.dropout if fn.train else 0
if (dropout_desc_name not in fn.dropout_state) or (
fn.dropout_state[dropout_desc_name].get() is None):
fn.dropout_state[dropout_desc_name] = Unserializable(
cudnn.DropoutDescriptor(handle, dropout_p, fn.dropout_seed))
dropout_desc = fn.dropout_state[dropout_desc_name].get()
dropout_desc.set_dropout(dropout_p, fn.dropout_seed)
return cudnn.RNNDescriptor(handle, fn.hidden_size, fn.num_layers,
dropout_desc, fn.input_mode, fn.bidirectional,
fn.mode, fn.datatype)
def init_dropout_descriptor(fn, handle):
return cudnn.DropoutDescriptor(handle, fn.dropout, fn.dropout_seed)