Optimal way to interface a model for multiple inferences?

[jalvesz]

Hi,
I would like to ask for guidance on how could I improve an interface to calling a model.

The Fortran code is a FEM solver, and I'm calling the model to update fields on gauss points. The issue I'm facing is that in the current implementation, I'm doing a call to the inference on each point instead of assembling the inputs for all points and doing an inference on an array of Tensors.

The workflow I'm doing is more or less the following:
I load the model, and I built an interface for inference like this:

subroutine torchscript_eval(id,tensor_in,tensor_out)
    integer, parameter :: wp = real32
    ! -- External Variables
    integer  :: id !> id of the model
    real(wp), target :: tensor_in(:)
    real(wp), target :: tensor_out(:)
    ! -- Internal Variables
    !-------------------------------------------------
    associate( p => ts_models_list(id) )
        !> wrap tensors
        p%in_tensor(1) = torch_tensor_from_array(tensor_in,  p%tensor_layout, p%device)
        p%out_tensor   = torch_tensor_from_array(tensor_out, p%tensor_layout, p%device)
        !> Run model and Infer
        call torch_module_forward(p%model, p%in_tensor, p%n_inputs, p%out_tensor)
        call torch_tensor_delete(p%in_tensor(1))
        call torch_tensor_delete(p%out_tensor)
    end associate
 end subroutine

derived type to contain the model and tensors

type ts_dplmat !> A silly container for the models
     character(len=:), allocatable :: name
     integer            :: in_features
     integer            :: out_features
     integer(c_int)     :: device = torch_kCPU
     type(torch_module) :: model

     integer :: n_inputs = 1
     integer :: tensor_layout(1) = [1]
     type(torch_tensor) :: in_tensor(1)
     type(torch_tensor) :: out_tensor
end type
type(ts_dplmat), allocatable :: ts_models_list(:) !> Companion list container

This function is being called within a do loop for all the gauss points.

So I did notice that wrapping and destroying the tensors with torch_tensor_from_array/torch_tensor_delete looks quite bad for a call that should be done within such a do loop... But if I pre-wrap the tensors and call the delete at the end, I saw the memory growing wildly, so I put it back in. The issue is that this does run veeery slowly.

This is my initial design, I have to iterate on that but just wanted to have your feedback if possible on using the interface to get the best performance possible.

Thanks,

[jatkinson1000]

Hi @jalvesz

You might like to take a look at the code here which is a more 'practical' example (benchmarking of running inference on multiple columns of an atmospheric model).

We have a setup routine to import the model and create the tensors that is run once at the start of the program, and a destruction routine to clean up that is run at the end. This reduces the overheads required at each iteration when running the model.
We also see that the first iteration after reading in the model is typically a little slower as things are created in memory (see discussion #55).

The net in the example I linked above operates on a single 'atmospheric column' of the numeric model, but we pass in multiple columns as a batch input for inference, but you may already be doing this.

Let me know if that is useful, and if you have further questions.

[KingLLee]

Hi @jatkinson1000
I noticed that 'practical' example at the code here have two versions (Forpy and of course Ftorch). Recently, I'm trying to optimise the running speed when embedding AI model to GCM and found that AI model initialization is important for speed. Also, I wonder if Forpy and Ftorch have the same operating efficiency when embedded in GCM.

[jatkinson1000]

Hi @KingLLee Yes, we were comparing the two to check timings.
We find that the timings are similar between the two, with FTorch perhaps a little faster on most models.
The main advantages here in using FTorch are a significantly simplified interface, no dependence on the python runtime, ad the MIT license.

Regarding initialisation efficiency, as discussed in conversations here we advise running initialisation of the model just once for the entire program, not every timestep/call.
In my codes with CESM v2 this is done in the phys_init subroutine of CAM.

[KingLLee]

Hi @jatkinson1000

Thanks for your comment for the comparison of Forpy and Ftorch. I'm happy to see that you reveal FTorch perhaps has a little faster speed on most models which gives me passionate to use this toolkit.

I have done several experiments to figure out how to use this efficiently. I note that besides model initialization, the less I embed AI model to loops, the faster the GCM models run. Not because the AI models runs much slower in one loop, but for the AI models are skilled at matrix operations and I can give it more information to solve in one time. So the next procedure of mine to optimise the running speed of GCM is to reduce the loop that contains the AI model.

Furthermore, I found that when I use model_init to initialise AI model, I must use model_end to kill the model on time, or my GCM will get stuck in the error of out of memory.

[jatkinson1000]

Hi @KingLLee Yes, passing data to the net as batches will be more efficient than calling it on each atmospheric column etc.
Yes - you need to destroy the model at the end of the code, AND the tensors after you have used them.
We are working on more verbose documentation, but I suggest working through the examples and also looking here where we call inference multiple times in a loop.

[KingLLee]

Thanks @jatkinson1000
Yes, of course,Thank you for the efforts of your entire team. I'm doing some AI-GCM experiments which is quite interesting and may be the future of model development.

[jalvesz]

Hi @jatkinson1000 thank for the advise! so, what I see as the main difference is that instead of using "torch_tensor_from_array" you used "torch_tensor_from_blob" in the benchmark, right? but you also set and delete the pointers within the loop. I tried to simplify my test as follows using your saved_simplenet_model_cpu.pt test model:

Click to open

program main
 use iso_fortran_env, only: sp=> real32, dp => real64
 use ftorch 
 implicit none
 type(torch_module) :: model ! Generate an object to hold the Torch model
 integer, parameter :: wp = sp
 ! Set up Fortran data structures
 real(wp), target :: in_data(5)
 real(wp), target :: out_data(5)
 integer, parameter :: n_inputs = 1
 integer :: tensor_layout(1) = [1]

 ! Set up Torch data structures
 type(torch_tensor) :: in_tensor(1)
 type(torch_tensor) :: out_tensor

 real(dp), allocatable :: array_dp(:,:)
 integer :: i, j
 real(dp) :: time_start, time_finish
 !===========================================================================
 ! Initialize the model from a TorchScript (.pt) file.
 model = torch_module_load('saved_simplenet_model_cpu.pt')

 ! Create Torch input/output tensors from the above arrays
 !in_tensor(1) = torch_tensor_from_array(in_data, tensor_layout, torch_kCPU)
 !out_tensor = torch_tensor_from_array(out_data, tensor_layout, torch_kCPU)

 !> first dummy loop to mimic a heavy iterative inference loop
 allocate( array_dp(5,10000) , source = 0._dp )
 call CPU_TIME(time_start)
 do j = 1,100
 array_dp(1:5,1) = [0.0, 1.0, 2.0, 3.0, 4.0]
 do i = 2, 10000
     in_data = real( array_dp(1:5,i-1) )
     call torchscript_eval( in_data , out_data )
     array_dp(1:5,i) = dble( out_data ) - array_dp(1:5,i-1)
 end do
 end do
 call CPU_TIME(time_finish)
 print *, 'time 1: ',time_finish - time_start
 deallocate( array_dp )

 !> second dummy loop to mimic a heavy iterative inference loop
 allocate( array_dp(5,10000) , source = 0._dp )
 call CPU_TIME(time_start)
 do j = 1,100
 array_dp(1:5,1) = [0.0, 1.0, 2.0, 3.0, 4.0]
 do i = 2, 10000
     in_data = real( array_dp(1:5,i-1) )
     call torchscript_eval_blob( in_data , out_data )
     array_dp(1:5,i) = dble( out_data ) - array_dp(1:5,i-1)
 end do
 end do
 call CPU_TIME(time_finish)
 print *, 'time 2: ',time_finish - time_start
 deallocate( array_dp )
 ! Clean up
 call torch_module_delete(model)
 
 contains
 
 subroutine torchscript_eval(tensor_in,tensor_out)
 ! -- External Variables
 real(wp), target :: tensor_in(:)
 real(wp), target :: tensor_out(:)
 ! -- Internal Variables
 !-------------------------------------------------
 !> wrap tensors
 in_tensor(1) = torch_tensor_from_array(tensor_in,  tensor_layout, torch_kCPU)
 out_tensor   = torch_tensor_from_array(tensor_out, tensor_layout, torch_kCPU)
 !> Run model and Infer
 call torch_module_forward(model, in_tensor, n_inputs, out_tensor)
 call torch_tensor_delete(in_tensor(1))
 call torch_tensor_delete(out_tensor)
 end subroutine
 
 subroutine torchscript_eval_blob(tensor_in,tensor_out)
 ! -- External Variables
 real(wp), target :: tensor_in(:)
 real(wp), target :: tensor_out(:)
 ! -- Internal Variables
 integer(c_int) :: ndims, layout(1)
 integer(c_int64_t) :: tensor_shape(1)
 !-------------------------------------------------
 !> wrap tensors
 ndims = 1
 tensor_shape = [5]
 layout = [1]
 in_tensor(1) = torch_tensor_from_blob(c_loc(tensor_in) , ndims, tensor_shape, layout, torch_kFloat32, torch_kCPU)
 out_tensor   = torch_tensor_from_blob(c_loc(tensor_out), ndims, tensor_shape, layout, torch_kFloat32, torch_kCPU)
 !> Run model and Infer
 call torch_module_forward(model, in_tensor, n_inputs, out_tensor)
 call torch_tensor_delete(in_tensor(1))
 call torch_tensor_delete(out_tensor)
 end subroutine
end program

Measuring the times of the two loops (one using torch_tensor_from_array and the other using torch_tensor_from_blob) in the main program give the same time: around 25.5 seconds each.

I also tried to analyze with VTune, I saw large times and memory being wasted somewhere but it has trouble showing where as I'm linking a "RelWithDebugInfo" against the release binaries of libtorch.

My intuition tells me that, ideally, call torch_module_forward( ), should be able to infer on the data but avoid creating new tensors within. Which is not the case as I see the memory ramping-up:

here the tensor_in/tensor_out have no effet, and uncommented these lines in the body of the program:

 ! Create Torch input/output tensors from the above arrays
 in_tensor(1) = torch_tensor_from_array(in_data, tensor_layout, torch_kCPU)
 out_tensor = torch_tensor_from_array(out_data, tensor_layout, torch_kCPU)

I'm trying to look more closely as what is happening inside but for the moment it is not clear to me.

[jatkinson1000]

Ah, yes, we use torch_tensor_from_blob, but this was written before we added the torch_tensor_from_array which should now be used in preference.

I'm going to tag @ElliottKasoar as I know he saw some issues with memory recently and has been working on a fix, though I believe that was in relation to GPU.
I'll try and revisit this tomorrow with fresh eyes.

As a comment looking at the code I would definitely try and pass array_dp in to the net and process as a batch rather than looping over i and calling the net many times.

[jalvesz]

Thanks so much :)

As a comment looking at the code I would definitely try and pass array_dp in to the net and process as a batch rather than looping over i and calling the net many times.

I know this is the ideal scenario, right now I'm kind of stuck with only being able to integrate in this context to showcase feasibility. But I do have in mind that I should get there eventually :)

[ElliottKasoar]

I'm going to tag @ElliottKasoar as I know he saw some issues with memory recently and has been working on a fix, though I believe that was in relation to GPU. I'll try and revisit this tomorrow with fresh eyes.

Ah so the memory bug I encountered wasn't actually GPU specific, and would lead to a similar increase in memory use over time, but unfortunately I still don't think it helps in this case.

The issue I saw was a result of reusing a variable for two different input tensors within the same loop (e.g. here and here, which meant when we cleaned up the tensor using torch_tensor_delete, only the data from the second input was actually deleted.

(I did only encounter it while running on GPU, but that was due to some combination of testing more iterations and different memory limits. It also meant I was able to use CUDACachingAllocator::setMemoryFraction to cause failure much faster, which is useful if the issue (might) relate to data moved to GPU.)

I'll have a think about why torch_module_forward alone might be to blame, but nothing is immediately obvious.

[TomMelt]

@jalvesz do you have a mwe or is your code open source? If so, I could take a look at it to see what's going on with the memory.

[jalvesz]

Hi @TomMelt thanks for your help. I wrote a mwe in the second comment, it is in an expandable section where it says "Click to open", the model I used there is the 'saved_simplenet_model_cpu.pt' coming from the examples in FTorch.

[TomMelt]

Can you please confirm that you are running the latest version of ftorch? I cannot reproduce the behaviour you are seeing.

I took your original MWE and compiled and run using the following:

gfortran -Og -g -I/software/ftorch/gfortran/include/ftorch -L/software/ftorch/gfortran/lib -lftorch -c test.f90
gfortran -Og -g -I/software/ftorch/gfortran/include/ftorch -o test test.o -L/software/ftorch/gfortran/lib -lftorch

When I run I get the following:

$ ./test 
 time 1:    6.0982089999999998     
 time 2:    6.0713189999999999  

To check memory usage I used valgrind (Note this will run much slower - roughly 20 times slower on my machine i.e., 240 s instead of 12s)

valgrind --tool=massif --detailed-freq=100000 --time-unit=ms ./test

I visualize the output using massif-visualizer:

$ massif-visualizer massif.out.*

I do not see an increase in memory once the main loops start running. Here is a screenshot of my plot:

And to check I deliberately added a fake allocation (fake_allocation) subroutine (see Modified MWE for full source below) that allocates 4MB - 4GB of RAM (doubling every 3 seconds).

Running the fake allocation (Modified MWE) gives the following memory usage:

subroutine fake_allocation(i)
  integer, intent(in) :: i
  integer, allocatable :: array(:)
  integer :: size_array
  size_array = 1024**2*2**i
  allocate( array(size_array) )
  array = 10
  call sleep(3)
  deallocate ( array )
end subroutine

Original MWE

program main
 use iso_fortran_env, only: sp=> real32, dp => real64
 use ftorch 
 implicit none
 type(torch_module) :: model ! Generate an object to hold the Torch model
 integer, parameter :: wp = sp
 ! Set up Fortran data structures
 real(wp), target :: in_data(5)
 real(wp), target :: out_data(5)
 integer, parameter :: n_inputs = 1
 integer :: tensor_layout(1) = [1]

 ! Set up Torch data structures
 type(torch_tensor) :: in_tensor(1)
 type(torch_tensor) :: out_tensor

 real(dp), allocatable :: array_dp(:,:)
 integer :: i, j
 real(dp) :: time_start, time_finish
 !===========================================================================
 ! Initialize the model from a TorchScript (.pt) file.
 model = torch_module_load('saved_simplenet_model_cpu.pt')

 ! Create Torch input/output tensors from the above arrays
 !in_tensor(1) = torch_tensor_from_array(in_data, tensor_layout, torch_kCPU)
 !out_tensor = torch_tensor_from_array(out_data, tensor_layout, torch_kCPU)

 !> first dummy loop to mimic a heavy iterative inference loop
 allocate( array_dp(5,10000) , source = 0._dp )
 call CPU_TIME(time_start)
 do j = 1,100
 array_dp(1:5,1) = [0.0, 1.0, 2.0, 3.0, 4.0]
 do i = 2, 10000
     in_data = real( array_dp(1:5,i-1) )
     call torchscript_eval( in_data , out_data )
     array_dp(1:5,i) = dble( out_data ) - array_dp(1:5,i-1)
 end do
 end do
 call CPU_TIME(time_finish)
 print *, 'time 1: ',time_finish - time_start
 deallocate( array_dp )

 !> second dummy loop to mimic a heavy iterative inference loop
 allocate( array_dp(5,10000) , source = 0._dp )
 call CPU_TIME(time_start)
 do j = 1,100
 array_dp(1:5,1) = [0.0, 1.0, 2.0, 3.0, 4.0]
 do i = 2, 10000
     in_data = real( array_dp(1:5,i-1) )
     call torchscript_eval_blob( in_data , out_data )
     array_dp(1:5,i) = dble( out_data ) - array_dp(1:5,i-1)
 end do
 end do
 call CPU_TIME(time_finish)
 print *, 'time 2: ',time_finish - time_start
 deallocate( array_dp )
 ! Clean up
 call torch_module_delete(model)
 
 contains
 
 subroutine torchscript_eval(tensor_in,tensor_out)
 ! -- External Variables
 real(wp), target :: tensor_in(:)
 real(wp), target :: tensor_out(:)
 ! -- Internal Variables
 !-------------------------------------------------
 !> wrap tensors
 in_tensor(1) = torch_tensor_from_array(tensor_in,  tensor_layout, torch_kCPU)
 out_tensor   = torch_tensor_from_array(tensor_out, tensor_layout, torch_kCPU)
 !> Run model and Infer
 call torch_module_forward(model, in_tensor, n_inputs, out_tensor)
 call torch_tensor_delete(in_tensor(1))
 call torch_tensor_delete(out_tensor)
 end subroutine
 
 subroutine torchscript_eval_blob(tensor_in,tensor_out)
 ! -- External Variables
 real(wp), target :: tensor_in(:)
 real(wp), target :: tensor_out(:)
 ! -- Internal Variables
 integer(c_int) :: ndims, layout(1)
 integer(c_int64_t) :: tensor_shape(1)
 !-------------------------------------------------
 !> wrap tensors
 ndims = 1
 tensor_shape = [5]
 layout = [1]
 in_tensor(1) = torch_tensor_from_blob(c_loc(tensor_in) , ndims, tensor_shape, layout, torch_kFloat32, torch_kCPU)
 out_tensor   = torch_tensor_from_blob(c_loc(tensor_out), ndims, tensor_shape, layout, torch_kFloat32, torch_kCPU)
 !> Run model and Infer
 call torch_module_forward(model, in_tensor, n_inputs, out_tensor)
 call torch_tensor_delete(in_tensor(1))
 call torch_tensor_delete(out_tensor)
 end subroutine
end program

Modified MWE

program main
  use iso_fortran_env, only: sp=> real32, dp => real64
  use iso_c_binding, only: c_associated, c_loc, c_ptr, c_int, c_int64_t
  use ftorch
  implicit none
  type(torch_module) :: model ! Generate an object to hold the Torch model
  integer, parameter :: wp = sp
  ! Set up Fortran data structures
  real(wp), target :: in_data(5)
  real(wp), target :: out_data(5)
  integer, parameter :: n_inputs = 1
  integer :: tensor_layout(1) = [1]

  ! Set up Torch data structures
  type(torch_tensor) :: in_tensor(1)
  type(torch_tensor) :: out_tensor

  real(dp), allocatable :: array_dp(:,:)
  integer :: i, j
  real(dp) :: time_start, time_finish
  !===========================================================================
  ! Initialize the model from a TorchScript (.pt) file.
  model = torch_module_load('saved_simplenet_model_cpu.pt')

  ! Create Torch input/output tensors from the above arrays
  !in_tensor(1) = torch_tensor_from_array(in_data, tensor_layout, torch_kCPU)
  !out_tensor = torch_tensor_from_array(out_data, tensor_layout, torch_kCPU)

  !> first dummy loop to mimic a heavy iterative inference loop
  allocate( array_dp(5,10000) , source = 0._dp )
  call CPU_TIME(time_start)
  do j = 1,100
    array_dp(1:5,1) = [0.0, 1.0, 2.0, 3.0, 4.0]
    do i = 2, 10000
      in_data = real( array_dp(1:5,i-1) )
      call torchscript_eval( in_data , out_data )
      array_dp(1:5,i) = dble( out_data ) - array_dp(1:5,i-1)
    end do
  end do
  call CPU_TIME(time_finish)
  print *, 'time 1: ',time_finish - time_start
  deallocate( array_dp )

  !> second dummy loop to mimic a heavy iterative inference loop
  allocate( array_dp(5,10000) , source = 0._dp )
  call CPU_TIME(time_start)
  do j = 1,100
    array_dp(1:5,1) = [0.0, 1.0, 2.0, 3.0, 4.0]
    do i = 2, 10000
      in_data = real( array_dp(1:5,i-1) )
      call torchscript_eval_blob( in_data , out_data )
      array_dp(1:5,i) = dble( out_data ) - array_dp(1:5,i-1)
    end do
  end do
  call CPU_TIME(time_finish)
  print *, 'time 2: ',time_finish - time_start
  deallocate( array_dp )
  ! Clean up
  call torch_module_delete(model)

  print *, 'fake allocation'
  do i = 1,10
    call fake_allocation(i)
  end do

contains

  subroutine fake_allocation(i)
    integer, intent(in) :: i
    integer, allocatable :: array(:)
    integer :: size_array
    size_array = 1024**2*2**i
    allocate( array(size_array) )
    array = 10
    call sleep(3)
    deallocate ( array )
  end subroutine

  subroutine torchscript_eval(tensor_in,tensor_out)
    ! -- External Variables
    real(wp), target :: tensor_in(:)
    real(wp), target :: tensor_out(:)
    ! -- Internal Variables
    !-------------------------------------------------
    !> wrap tensors
    in_tensor(1) = torch_tensor_from_array(tensor_in,  tensor_layout, torch_kCPU)
    out_tensor   = torch_tensor_from_array(tensor_out, tensor_layout, torch_kCPU)
    !> Run model and Infer
    call torch_module_forward(model, in_tensor, n_inputs, out_tensor)
    call torch_tensor_delete(in_tensor(1))
    call torch_tensor_delete(out_tensor)
  end subroutine

  subroutine torchscript_eval_blob(tensor_in,tensor_out)
    ! -- External Variables
    real(wp), target :: tensor_in(:)
    real(wp), target :: tensor_out(:)
    ! -- Internal Variables
    integer(c_int) :: ndims, layout(1)
    integer(c_int64_t) :: tensor_shape(1)
    !-------------------------------------------------
    !> wrap tensors
    ndims = 1
    tensor_shape = [5]
    layout = [1]
    in_tensor(1) = torch_tensor_from_blob(c_loc(tensor_in) , ndims, tensor_shape, layout, torch_kFloat32, torch_kCPU)
    out_tensor   = torch_tensor_from_blob(c_loc(tensor_out), ndims, tensor_shape, layout, torch_kFloat32, torch_kCPU)
    !> Run model and Infer
    call torch_module_forward(model, in_tensor, n_inputs, out_tensor)
    call torch_tensor_delete(in_tensor(1))
    call torch_tensor_delete(out_tensor)
  end subroutine
end program

[jatkinson1000]

Thanks @TomMelt
@jalvesz hopefully this is useful and you can re-check your results with the latest version of FTorch.

[jalvesz]

Hi @TomMelt,
Regarding the version, I'm indeed using the latest one. With the test as is, indeed there is no memory build up, but what I found is that the program runs slowly. So, I tried to see what happens if I just call the forward evaluation and the pointers setting/destruction outside the loop:

  subroutine torchscript_eval(tensor_in,tensor_out)
    ! -- External Variables
    real(wp), target :: tensor_in(:)
    real(wp), target :: tensor_out(:)
    ! -- Internal Variables
    !-------------------------------------------------
    !> Run model and Infer
    call torch_module_forward(model, in_tensor, n_inputs, out_tensor)
  end subroutine

Here is where I saw the memory build-up.

[TomMelt]

Hi could you please send me the exact code that demonstrates memory build up. I am not 100% I know what you have done. I would like to investigate it further.

[jatkinson1000]

Hi @jalvesz we haven't heard anything and are unable to reproduce, so I'm closing this as outdated.
Please do let us know if you want to re-open this however.