How to properly make batch inference?

[IsidreMas]

Hello FTorch developers,

I'm a new user of the library trying to include it in an atmospheric chemistry emulator to interface with ML models. So far I managed to reproduce the inference results running the model 1 sample at a time, but when attempting to do batch inference the results are not numerically correct. Do you know what I could be doing wrong?

Here's the example code I'm using to load a dataset from a netcdf file and write the predicitons of the model to a file:

batchintest.F90

program read_toy
    use, intrinsic :: iso_fortran_env, only: sp => real32, int64, dp => real64
    use netcdf
    ! IMPORT TORCH INTERFACE MODULES
    use ftorch, only: torch_model, torch_tensor, &
         torch_kCPU, torch_kCUDA, torch_kXPU, torch_kMPS, &
         torch_tensor_from_array, torch_model_load, torch_model_forward, &
         torch_delete
    implicit none

    ! SET PRECISION
    integer, parameter :: wp = sp

    integer :: ncid, status, varid, i, num_samples
    real(wp), allocatable :: gas_conc(:,:,:) ! 3D array to store gas concentrations

    ! COMMAND-LINE ARGUMENTS
    integer :: num_args, ix, device_index
    character(len=128), dimension(:), allocatable :: args

    ! USER CONFIGURATION: CHANGE INPUT/OUTPUT DIMENSIONS AS NEEDED.
    integer, parameter :: n_features = 3      ! (E.G., FOR CALIFORNIA HOUSING, THIS IS 8)
    integer, parameter :: n_outputs  = 3        ! (CHANGE IF YOUR MODEL RETURNS MULTIPLE VALUES)

    ! INPUT DATA ARRAY (SAMPLE): MODIFY THESE VALUES FOR YOUR USE-CASE.
    real(wp), dimension(:,:), target, allocatable :: in_data
    ! OUTPUT DATA ARRAY (INITIALLY ZEROED)
    real(wp), dimension(:,:), target, allocatable :: out_data

    ! DEFINE TENSOR LAYOUT: BATCH SIZE = 1 (SINGLE SAMPLE).
    integer, parameter :: tensor_layout(2) = [1,2]

    ! TORCH OBJECTS
    type(torch_model) :: model
    type(torch_tensor), dimension(:), allocatable :: in_tensors
    type(torch_tensor), dimension(:), allocatable :: out_tensors

    ! DEVICE SELECTION VARIABLES
    integer :: device_type, num_devices

    ! FILE I/O UNIT NUMBER FOR SAVING PREDICTIONS
    integer, parameter :: outfile_unit = 10

    ! TIMING VARIABLES
    integer(int64) :: start_time, end_time, count_rate
    real(dp) :: elapsed_time, total_time = 0.0_dp

    ! PARSE COMMAND-LINE ARGUMENTS
    num_args = command_argument_count()
    allocate(args(num_args))
    do ix = 1, num_args
        call get_command_argument(ix, args(ix))
    end do

    ! DEVICE SELECTION BASED ON FIRST ARGUMENT (REQUIRED)
    if (trim(args(1)) == "cuda") then
        device_type = torch_kCUDA
        num_devices = 2
    else if (trim(args(1)) == "xpu") then
        device_type = torch_kXPU
        num_devices = 2
    else if (trim(args(1)) == "mps") then
        device_type = torch_kMPS
        num_devices = 1
    else if (trim(args(1)) == "cpu") then
        device_type = torch_kCPU
        num_devices = 1
    else
        write (*,*) "Error: invalid device type", trim(args(1))
        stop 999
    end if

    ! Open NetCDF file
    status = nf90_open("toy.nc", NF90_NOWRITE, ncid)
    if (status /= nf90_noerr) call handle_err(status)

    ! Read gas concentrations (3D: species × experiment × time)
    status = nf90_inq_varid(ncid, "gas_species_concentrations", varid)
    call handle_err(status)
    allocate(gas_conc(3, 1010688, 2))
    status = nf90_get_var(ncid, varid, gas_conc)
    print *, "First gas concentration:", gas_conc(:,1,1)

    num_samples = size(gas_conc, 2)

    allocate(in_tensors(1))
    allocate(in_data(num_samples, n_features))
    allocate(out_tensors(1))
    allocate(out_data(num_samples, n_outputs))

    ! LOAD THE TORCHSCRIPT MODEL FROM THE FILE (SECOND ARGUMENT)
    call torch_model_load(model, args(2), device_type)

    ! Transpose gas concentrations to match the input tensor shape
    in_data = transpose(gas_conc(:, :, 1))
    
    print *, "First few rows of gas_conc:"
    do i = 1, 5
        print *, gas_conc(:, i, 1)
    end do

    ! Get the system clock rate
    call system_clock(count_rate=count_rate)

    print *, "First few rows of input data:"
    do i = 1, 5
        print *, in_data(i, :)
    end do

    ! CREATE THE TORCH INPUT TENSOR ON THE SELECTED DEVICE
    call torch_tensor_from_array(in_tensors(1), in_data, tensor_layout, device_type)

    ! CREATE THE TORCH OUTPUT TENSOR ON THE CPU
    call torch_tensor_from_array(out_tensors(1), out_data, tensor_layout, torch_kCPU)

    print *, "First input tensor:", in_tensors(1)%get_shape()
    print *, "First output tensor:", out_tensors(1)%get_shape()
    print *, "in_data shape:", shape(in_data)
    print *, "out_data shape:", shape(out_data)

    ! TIME THE FORWARD PASS
    call system_clock(start_time)
    call torch_model_forward(model, in_tensors, out_tensors)
    call system_clock(end_time)

    ! Calculate elapsed time in seconds
    elapsed_time = real(end_time - start_time, dp) / real(count_rate, dp)
    total_time = total_time + elapsed_time

    ! Print total prediction time
    print '(a, f12.6, a)', 'Total prediction time: ', total_time, ' seconds'

    ! Open predictions.txt file for writing (overwrite if exists)
    open(unit=outfile_unit, file="predictions.txt", action="write", status="replace")
    ! WRITE THE OUTPUT DATA TO predictions.txt
    write(outfile_unit, '(3(F20.10))') out_data

    ! CLEANUP: DELETE THE MODEL AND TENSORS
    call torch_delete(model)
    call torch_delete(in_tensors)
    call torch_delete(out_tensors)

    ! Close the output file
    close(outfile_unit)

    ! Cleanup
    deallocate(gas_conc)
    status = nf90_close(ncid)

contains
    subroutine handle_err(status)
        integer, intent(in) :: status
        if (status /= nf90_noerr) then
            print *, "Error: ", trim(nf90_strerror(status))
            stop
        end if
    end subroutine
end program read_toy

Thank you for any hint or advice that you could provide on this issue.

[jatkinson1000]

Hi @IsidreMas At the moment batching works easily in Python because of the array broadcasting capabilities (you can pass an 1 x n and an m x n tensor into a net that expects 1 x n and it behaves accordingly).

In FTorch this is not currently the case as there is no built-in batching, so you have to loop over your m x n tensor m times, running the net over each 1 x n vector/column.

We would like to make the interface similar to PyTorch's and allow a range of shapes to be passed as batches, but this is not something we have currently looked at in-depth. You can see an open issue for this feature at #157

We would be happy to work with you on implementing this, and have developer guidelines.
The naiive approach would loop over the vectors in the Fortran/C++, but for efficiency I think we want to try and let this happen at the LibTorch level to make use of underlying optimizations.
With a quick look something like the following in C++ would work:

c++ batching in LibTorch

#include <torch/torch.h>
#include <torch/script.h>

int main() {
  // Load the TorchScript model
  torch::jit::Module model;
  model = torch::jit::load("model.pt");

  // Create a batched input tensor (m × n)
  int64_t batch_size = 5; // m
  int64_t input_size = 4; // n
  torch::Tensor input = torch::rand({batch_size, input_size});

  // Forward pass through the model
  torch::Tensor output;
  output = model.forward({input}).toTensor();

  // Print the output
  std::cout << "Output:\n" << output << std::endl;

  return 0;
}

Which indicates batching should be built into LibTorch so doable.
I wonder if care needs to be taken over the tensor_layout as LibTorch always expects the batch index to be the first dimension of the tensor.
I see you do this in your code, but would need to review exactly what happens when we move from row to column major in the underlying FTorch.

What would be really useful is if you could perhaps write a minimal working example of FTorch code - perhaps a modification of one of the examples we provide. We can then use this to look at how mest to implement batching and make progress with you.

[IsidreMas]

Hi @jatkinson1000 , thank you for the reply, indeed it would be nice to have the batching capability of libtorch seamlessly integrated. I think the current implementation might not be so far away from it as it might seem, actually the code I first shared was working but the print of the output tensor was not done properly. Specifically,

! WRITE THE OUTPUT DATA TO predictions.txt
  write(outfile_unit, '(3(F20.10))') out_data

doesn't print the results in the correct order, but if I print the results one by one it seems like it worked:

! WRITE THE OUTPUT DATA TO predictions.txt
do i = 1, num_samples
    write(outfile_unit, '(3(F20.10))') out_data(i, :)
end do

The forward pass took 0,26 seconds to process 1010688 samples, while it took 16 seconds to process them in a m times loop going in samples of 1xn.

I can try to provide a modification of one of the examples with this batch processing working if that helps.

Comparison results

This are some comparison results between batch processing on python vs fortran:

Batch size: 1010688
Running inference on CPU...
CPU inference time: 0.0408 seconds
Running inference on GPU...
GPU inference time: 0.0022 seconds
Mean absolute difference between CPU and GPU outputs on Python: 1.843717e-09
First few inferred predictions (PYTHON):
[[ 0.03851714 -0.03406327 0.01425657]
[ 0.03775792 -0.03478099 0.01411733]
[ 0.03966463 -0.02383732 -0.0056127 ]
[ 0.02826056 -0.00427802 0.00210775]
[ 0.0308076 -0.00167534 -0.00681062]]

In fortran it took 0,16 s to batch process on CPU
First few file predictions (FORTRAN) :
[[ 0.03851714 -0.03406327 0.01425657]
[ 0.03775792 -0.03478099 0.01411733]
[ 0.03966463 -0.02383732 -0.0056127 ]
[ 0.02826056 -0.00427802 0.00210775]
[ 0.0308076 -0.00167534 -0.00681062]]
Mean absolute difference between inferred predictions and file predictions: 2.202213e-09

It seems however that the batch processing is still faster in python (0.0408 s) than in Fortran (0.16 s), maybe there's some room for optimization there or perhaps I'm not measuring the time correctly