Torch interoperability

Send data to torch

When using ultraspy on GPU, the objects are sent to GPU using the cupy library, which is dealing with the pointers to know where to find the data within the GPU device.

If you aim to use some Deep Learning algorithm, or simply to run your own algorithm using the torch library, it is possible to send the information to the torch library. You can do it the following way:

 1import numpy as np
 2import cupy as cp
 3import torch
 4
 5# Some random data
 6a = np.random.randint(10, size=(5, 6))
 7
 8# Send it to device using cupy, as ultraspy does
 9cupy_a = cp.array(a, np.float32)
10
11# Transfer its pointer to torch so it can be used in both libraries
12torch_a = torch.as_tensor(cupy_a, device='cuda')

Now, the array a is available both on CPU (the numpy array called a) and on GPU (accessible both by cupy with cupy_a and torch with torch_a). Note that both cupy_a and torch_a refers to the same data on GPU, meaning that if you update one, the other one will also be updated. Let’s say for example that you want to convert the data as decibel using both libraries you can do:

 1# Normalize the data using cupy
 2cupy_a /= cupy_a.max()
 3
 4# Convert to log using torch
 5torch.log10(torch_a, out=torch_a)
 6
 7# Get the data back on CPU
 8cpu_cupy_a = cupy_a.get()
 9cpu_torch_a = torch_a.cpu()
10
11# Both are equal
12print(np.allclose(cpu_cupy_a, cpu_torch_a))

All good. Few things to consider here though. It works as long as we are modifying directly the device array. Make sure that all the operations are inplace, and won’t allocate a new array in memory in the process. If a new array is allocated by either cupy or torch, the other lib won’t be aware of the update:

1# Operation is modifying both libraries
2cupy_a /= cupy_a.max()
3print(np.allclose(cupy_a.get(), torch_a.cpu()))  # True
4
5# The array in cupy is divided by max, then stored in a new array on GPU,
6# that happens to have the same name, torch and cupy are now different
7# pointers
8cupy_a = cupy_a / cupy_a.max()
9print(np.allclose(cupy_a.get(), torch_a.cpu()))  # False

Similarly in torch:

1# The log10, by default, is allocating a new array, but we can specify the
2# output location. If it is the same as before, it'll still be the same
3# pointer as cupy
4torch.log10(torch_a, out=torch_a)
5print(np.allclose(cupy_a.get(), torch_a.cpu()))  # True
6
7# A new pointer is created, cupy array is detached from torch'
8torch_a = torch.log10(torch_a)
9print(np.allclose(cupy_a.get(), torch_a.cpu()))  # False

That’s it! And if you want to go the other way, and converting your torch Tensor to a cupy array, you should use:

1a = np.random.randint(10, size=(5, 6))
2torch_a = torch.as_tensor(a)
3cupy_a = cp.asarray(torch_a)

Load a Deep Learning model

A good way to call a Deep Learning model from ultraspy could be to embed it within a Singleton class. That way, we can ensure that the model is instantiated only once. Any other way would work too based on your application, but we’ll focus on the Singleton in this example.

Let’s imagine we’ve trained a network to process three tilted plane waves to produce results as good as if we had sent 31 plane waves. A way to embed it could be:

 1class MyModelSingleton:
 2    # General information
 3    _instance = None
 4    _model = None
 5    _weights = "torch_models/my_model_weights.pth"
 6
 7    def __new__(cls, *args, **kwargs):
 8        # Initializer, only once
 9        if cls._instance is None:
10            cls._model = MyModelClass().to("cuda")
11            cls._model.load_state_dict(torch.load(cls._weights))
12            cls._model.eval()
13            cls._instance = super().__new__(cls)
14        return cls._instance
15
16    def prepare_data(self, d_data):
17        # Prepare the data, and transfer its pointer to torch
18        return torch.as_tensor(d_data, device='cuda')
19
20    def run_inference(self, t_data):
21        # Runs the inference and return the prediction
22        with torch.inference_mode():
23            prediction = self._model(t_data)
24        return prediction
25
26    def back_to_cupy(self, t_data):
27        # Returns the data to the cupy lib
28        return cp.asarray(t_data)

That way, our function would be very simple:

1def deep_compound(d_beamformed):
2    # The model is ready for use after the first instantiation
3    model = CidNetSingleton()
4
5    input_model_tensor = model.prepare_data(d_beamformed)
6    prediction = model.run_inference(input_model_tensor)
7    return model.back_to_cupy(prediction)