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()

nnabla.ext_utils.get_extension_context(ext_name, **kw)[ソース]

指定されたExtensionのContextを取得します。

すべてのExtensionモジュールは context(**kw) 関数を提供する必要があります。

パラメータ:

ext_name (str) -- nnabla_ext への相対的なモジュールパス
kw (dict) -- Extensionモジュールにあるcontext関数への追加のキーワード引数。

戻り値:

現在のExtensionのcontext

戻り値の型:

nnabla.Context

サンプル

ctx = get_extension_context('cudnn', device_id='0', type_config='half')
nn.set_default_context(ctx)

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 を取りますが、これによって同期させたいデバイスを指定できます。