### カスタムプラグイン実装の意義
TensorRTの利用中に算子非対応や生成構造冗長化、fp16/int8精度誤差増大などの問題が発生する場合があります。こうした課題に対しては独自算子の実装が必要です。onnx-graphsurgeonによりONNXモデルの編集(ノード追加・削除・接続)を行い、Polygraphyを併用して精度検証を実施します。onnx-graphsurgeonはNVIDIAが提供するONNX編集ツール、Polygraphyはモデル前後比較と逐層精度チェックに対応する同社ツールです。
### ONNXプラグインを含むモデル構築
カスタムプラグイン構築にはトレーニング時モデル構築、ONNXエクスポート、および推論プラットフォーム変換の3段階があります。以下はSwish演算モジュールのPyTorch実装例とTensorRTプラグイン登録プロセスです。
/savedemo/Plugin_Examplesディレクトリにデモコードが格納されています。
PyTorchでのSwish実装には前向伝播・逆伝播・ONNXエクスポートの3要素が必要です。
class MySwish(torch.autograd.Function):
@staticmethod
def symbolic(g, input, bias):
return g.op("custom::Plugin", input, bias, name_s="Swish", info_s=json.dumps({
"size": 555,
"shape": [6, 7, 8],
"module": "abcdefg"
}))
@staticmethod
def forward(ctx, input, bias):
ctx.save_for_backward(input)
return input * torch.sigmoid(input) + bias
@staticmethod
def backward(ctx, grad_output):
input = ctx.saved_tensors[0]
sigmoid_input = torch.sigmoid(input)
return grad_output * (sigmoid_input * (1 - sigmoid_input)), grad_output
勾配チェックは以下の方法で実行可能です。
input_tensor = torch.ones(3, requires_grad=True, dtype=torch.double)
bias = torch.ones(3, requires_grad=True, dtype=torch.double)
input_check = gradcheck(MySwish.apply, (input_tensor, bias), eps=1e-5, atol=1e-4)
bias_check = gradcheck(MySwish.apply, (input_tensor, bias), eps=1e-5, atol=1e-4)
print("入力勾配チェック:", input_check)
print("バイアス勾配チェック:", bias_check)
出力がTrueのとき正しく動作しています。
モデル構築ではカスタム算子を以下の形式で使用します。
class EfficientSwish(nn.Module):
def __init__(self):
super().__init__()
self.bias = nn.Parameter(torch.zeros(1))
self.bias.data.fill_(3.15)
def forward(self, x):
return MySwish.apply(x, self.bias)
3×3畳み込みネットワークへの組み込み例:
class Net(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(1, 1, 3, stride=1, padding=0, dilation=1, bias=False)
self.conv.weight.data = torch.FloatTensor([
[1, 0, 0],
[0, 1, 0],
[0, 0, 1]
]).view(1, 1, 3, 3)
self.swish = EfficientSwish()
def forward(self, x):
return self.swish(self.conv(x))
ONNXエクスポート時には注意点が2つあります。1つはドメイン名の指定形式(scope::plugin)、もう1つはopsetバージョンの設定です。
domain = "custom"
torch.onnx.register_custom_op_symbolic(f"{domain}::Plugin", MySwish.symbolic, 1)
torch.onnx.export(model, (input,), "plugin.onnx",
input_names=["Input"],
output_names=["Output"],
opset_version=11,
verbose=True)
### TensorRTプラグイン実装
ONNXエクスポート後のTensorRTデプロイではC++インターフェースを使用します。主要なステップは以下の通りです:1.プラグイン実装 2.プラグイン作成クラス 3.TRTRへの登録 4.シリアライズ処理。デシリアライズプロセスではプラグインタイプ・バージョン取得→クリエイター取得→ネットワーク挿入の流れとなります。
プラグインコンストラクタの基本構造:
// ネットワーク定義時・クローン時・クリエイター呼び出し時に呼ばれます
// 定義時のコンストラクタは手動で実装が必要
// クローンはオプショナルです
// 定義時コンストラクタのパラメータは柔軟に設定可能
// デシリアライズ時はデータアドレスと長さのみ受け取ります
主要API仕様:
int getNbOutputs() const; // 出力数取得
nvinfer1::Dims getOutputDimensions(int index, const nvinfer1::Dims* inputs, int nbInputDims); // 出力寸法取得
nvinfer1::DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const; // 出力データ型取得
size_t getSerializationSize() const; // パラメータサイズ取得
void serialize(void* buffer) const; // パラメータ書き込み
const char* getPluginType() const; // プラグイン名取得(ユニーク)
const char* getPluginVersion() const; // バージョン取得(ユニーク)
int initialize(); // 重みメモリ確保
void terminate(); // メモリ解放
void destroy(); // プラグイン全体解放
void configurePlugin(const nvinfer1::PluginTensorDesc* in, int nbInput, const nvinfer1::PluginTensorDesc* out, int nbOutput); // 初期化
bool supportsFormatCombination(int pos, const nvinfer1::PluginTensorDesc* inOut, int bnInputs, int bnOutputs) const; // フォーマット検証
size_t getWorkspaceSize(int maxBatchSize) const; // ワークスペース取得
int enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream); // 実行
nvinfer1::IPluginV2* deserializePlugin(const char* name, const void* serialData, size_t serialLength); // デシリアライズ
前処理/後処理モジュールの場合、PyTorch実装は不要でONNXノード追加後にTensorRTで直接実装可能です。ONNXノード操作方法については別途ドキュメントをご参照ください。ONNXノードにはop(操作タイプ)、name(ノード名)、attrs(属性辞書)、inputs/outputs(入出力情報)が含まれます。Tensorはname、dtype、shapeの3属性を持ち、shapeは入力ノードのみ事前確定です。重みを含む場合はattrsに情報を追記します。
詳細な実装例は/savedemo/Plugin_Examplesディレクトリにコメント付きで収録されています。