Extension¶
NNabla offers easy extensibility for developers to add new device extensions.
The NNabla Python package officially supports the cpu
, cuda
and cudnn
extension,
cuda
and cudnn
extension can dramatically accelerate computation by leveraging NVIDIA
CUDA GPUs with cuDNN computation primitives.
以下のように手動でExtensionをインポートできます:
import nnabla_ext.cudnn
See :ref:`python-package-installation` to install the CUDA extension.
Extensionユーティリティ¶
NNabla Extensionのためのユーティリティ。
- nnabla.ext_utils.list_extensions()[ソース]¶
利用可能なExtensionをリストアップします。
注釈
名前空間パッケージのディレクトリ構成に依存しているため、あるプラットフォームや環境によっては正しく動作しない可能性があります。
- 戻り値: str のリスト
利用可能なExtension名のリスト。
- nnabla.ext_utils.import_extension_module(ext_name)[ソース]¶
Extensionモジュールの名称でインポートします。
Extensionモジュールは、
nnabla_ext
パッケージの下に名前空間パッケージとしてインストールされます。 すべてのExtensionモジュールは、統一した API セットを提供します。- パラメータ
ext_name (str) -- Extension名。例: ‘cpu’, ‘cuda’, ‘cudnn’ 等。
- 戻り値: module
特定の NNabla Extensionの Python モジュール。
サンプル
ext = import_extension_module('cudnn') available_devices = ext.get_devices() print(available_devices) ext.device_synchronize(available_devices[0]) ext.clear_memory_cache()
Extensionモジュールの API¶
すべてのExtensionモジュールは以下の関数を持つ必要があります。
- nnabla.ext_utils.context(*kw)¶
Extensionモジュールのデフォルトのcontext記述子を返します。本メソッドはExtensionによって異なるオプションとなる引数を取りえます。例えば、
cudnn
Extensionでは、計算処理を実行する GPU を指定するためint
としてdevice_id
を取ります。
- nnabla.ext_utils.device_synchronize(*kw)¶
ホストスレッドに対してデバイス実行ストリームを同期するために本メソッドを使います。例えば CUDA では、カーネル実行をストリームのキューに追加し、ホストスレッドとは非同期に実行します。本関数は、そのような機能を使うデバイス内でのみ有効です。CPU 実装では、本メソッドはダミー関数として実装され、そのため本関数を呼び出しても無視されます。
cudnn
Extensionにある本関数は、オプションの引数としてdevice_id
を取りますが、これによって同期させたいデバイスを指定できます。