zea.backend.torch

Pytorch Ultrasound Beamforming Library.

Initialize modules for registries.

Functions

on_device_torch(func, inputs, device[, ...])

Applies a Torch function to inputs on a specified device.
zea.backend.torch.on_device_torch(func, inputs, device, return_numpy=False, **kwargs)[source]

Applies a Torch function to inputs on a specified device.

Parameters:

  • func (function) – Function to apply to the input data.

  • inputs (ndarray) – Input array.

  • device (str) – Device string, e.g. 'cuda', 'gpu', or 'cpu'.

  • return_numpy (bool, optional) – Whether to convert output data back to numpy. Defaults to False.

  • **kwargs – Additional keyword arguments to be passed to the func.

Returns:

The output data.

Return type:

torch.Tensor or ndarray

Raises:

AssertionError – If func is not a function from the torch library.

Note

This function converts the inputs array to a torch.Tensor and moves it to the specified device. It then applies the func function to the inputs and returns the output data. If the output is a dictionary, it extracts the first value from the dictionary. If return_numpy is True, it converts the output data back to a numpy array before returning.

Example

import torch


def square(x):
    return x**2


inputs = [1, 2, 3, 4, 5]
device = "cuda"
output = on_device_torch(square, inputs, device)
print(output)

Modules

losses

Container for custom loss functions.