Accessing Hidden State Using C API

[vectorphresh]

Hello,

I wanted to know anyone has had success accessing a models hidden state using C Training API. I thought that I could export to ONNX with the hidden state as both an input and output parameter of the model, but when I run the TrainStep function my program crashes. The documentation hints that I may be able to access model parameters directly, but I'm not quiet getting it. Here's my training function. I instantiated the hidden layer outside of this function:

Update 10/14/2024: I realized trying to init the tensors in a separate function may have not been using the same memory allocator, I instead am saving the floating point representation of hidden layer. Still no dice on completing call to TrainStep. Next step is try compiling with debug.

int model_run_training_step(model_context_t *ctx, float target_value, int target_size) {
    OrtValue *input_tensor = NULL;
    OrtValue *output_hidden_tensor = NULL;
    OrtValue *target_tensor = NULL;
    float *output_hidden_state = NULL;
    OrtMemoryInfo *memory_info;

    log_debug("Allocating memory for inference...");
    if (ctx->cuda_enabled && strcmp(ctx->info->memory_type, "Cuda") == 0) {
        ORT_ABORT_ON_ERROR(ctx->ort->CreateMemoryInfo(ctx->info->memory_type, OrtArenaAllocator, OrtMemTypeDefault, 0, &memory_info));
    } else {
        ORT_ABORT_ON_ERROR(ctx->ort->CreateCpuMemoryInfo(OrtArenaAllocator, OrtMemTypeDefault, &memory_info));
    }

    // Create input tensor
    float *input_data = min_max_scale_and_copy_array(ctx->info->features, ctx->info->sequence_length, ctx->info->num_features);
    if (!input_data) {
        log_error("Failed to allocate input_data");
        return -1;
    }
    int input_size = ctx->info->sequence_length * ctx->info->num_features;
    log_debug("Input tensor shape: [%lld, %lld, %lld]", ctx->info->input_shape[0], ctx->info->input_shape[1], ctx->info->input_shape[2]);
    ORT_ABORT_ON_ERROR(ctx->ort->CreateTensorWithDataAsOrtValue(memory_info, input_data, sizeof(float) * input_size,
            ctx->info->input_shape, ctx->info->input_dims, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, &input_tensor));

    // Create target tensor
    int64_t target_dims[] = {1, target_size};
    log_debug("Target tensor shape: [%lld, %lld]", target_dims[0], target_dims[1]);
    ORT_ABORT_ON_ERROR(ctx->ort->CreateTensorWithDataAsOrtValue(memory_info, &target_value, sizeof(float) * target_size,
            target_dims, 2, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, &target_tensor));

    // Allocate hidden state if NULL
    if (ctx->info->hidden_state == NULL) {
        log_debug("Hidden state is NULL, allocating...");
        ctx->info->hidden_state = calloc(1, ctx->info->hidden_data_len);
        if (!ctx->info->hidden_state) {
            log_error("Failed to allocate hidden_state");
            return -1;
        }
    }

    // Create hidden state tensor
    log_debug("Hidden state tensor shape: [%lld, %lld, %lld]", ctx->info->hidden_shape[0], ctx->info->hidden_shape[1], ctx->info->hidden_shape[2]);
    ORT_ABORT_ON_ERROR(ctx->ort->CreateTensorWithDataAsOrtValue(memory_info, ctx->info->hidden_state,
            ctx->info->hidden_data_len, ctx->info->hidden_shape, ctx->info->hidden_dims, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, &ctx->hidden_state_tensor));

    // Allocate output hidden state
    output_hidden_state = calloc(1, ctx->info->hidden_data_len);
    if (!output_hidden_state) {
        log_error("Failed to allocate output_hidden_state");
        return -1;
    }
    ORT_ABORT_ON_ERROR(ctx->ort->CreateTensorWithDataAsOrtValue(memory_info, output_hidden_state,
            ctx->info->hidden_data_len, ctx->info->hidden_shape, ctx->info->hidden_dims, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, &output_hidden_tensor));

    pthread_mutex_lock(&onnx_mutex);

    // Training step
    log_debug("Performing Training Step...");
    ORT_ABORT_ON_ERROR(ctx->ort_train->TrainStep(ctx->train_session, NULL,
            2, (const OrtValue *const[]){input_tensor, ctx->hidden_state_tensor}, 2, (OrtValue *[]){&target_tensor, &output_hidden_tensor}));

    pthread_mutex_unlock(&onnx_mutex);

    // Apply optimizer step
    ORT_ABORT_ON_ERROR(ctx->ort_train->OptimizerStep(ctx->train_session, NULL));

    // Reset gradients
    ORT_ABORT_ON_ERROR(ctx->ort_train->LazyResetGrad(ctx->train_session));

    // Capture and reuse updated hidden state
    float *updated_hidden_state = NULL;
    log_debug("Retrieving updated hidden state...");
    ORT_ABORT_ON_ERROR(ctx->ort->GetTensorMutableData(output_hidden_tensor, (void **)&updated_hidden_state));
    memmove(ctx->info->hidden_state, updated_hidden_state, ctx->info->hidden_data_len * sizeof(float));

    // Cleanup
    ctx->ort->ReleaseValue(input_tensor);
    ctx->ort->ReleaseValue(target_tensor);
    ctx->ort->ReleaseValue(output_hidden_tensor);
    free(output_hidden_state);

    return 0;
}

OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 22.04
ONNX Runtime installed from (source or binary): binary
ONNX Runtime version: onnxruntime-training-linux-x64-1.16.3.tgz
ONNX Language: C
Visual Studio version (if applicable): Using Visual Studio Code
GCC/Compiler version (if compiling from source): 11.4.0-1ubuntu1~22.04
CUDA/cuDNN version: N/A (Cpu only. was trying to get Cpu working before exporting for cuda)
GPU model and memory: RTX 4080 16GB