APIs
PyTorch Native APIs¶
PyTorch의 대부분의 native 함수들은 RBLN NPU에서도 그대로 사용할 수 있습니다.
RBLN 특화 API¶
다음 함수들은 torch.rbln module에서 정의된 RBLN NPU 특화 함수들입니다. 명확한 이해를 위해 영문으로 작성되어 있습니다.
Classes¶
device_of
¶
Context-manager that changes the current device to that of given object.
You can use both tensors and storages as arguments. If a given object is not allocated on an RBLN device, this is a no-op.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
obj
|
Tensor or Storage
|
object allocated on the selected device. |
required |
Functions¶
get_amp_supported_dtype()
¶
Get a list of data types supported by automatic mixed precision (AMP) on RBLN devices.
Returns:
| Type | Description |
|---|---|
List[dtype]
|
List[torch.dtype]: A list of data types supported by AMP. |
Note
This function currently returns only torch.float16.
is_available()
¶
Check if any RBLN devices are available.
Returns:
| Name | Type | Description |
|---|---|---|
bool |
bool
|
True if at least one RBLN device is available, False otherwise. |
current_device()
¶
Get the index of the currently selected RBLN device.
Returns:
| Name | Type | Description |
|---|---|---|
int |
int
|
The index of the currently selected RBLN device. |
device_count()
¶
Get the number of available RBLN devices.
Returns:
| Name | Type | Description |
|---|---|---|
int |
int
|
The number of available RBLN devices. |
set_device(device)
¶
Set the current device.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
device
|
device or int or str
|
selected device. |
required |