深層学習カスタムプラグインの作成

### カスタムプラグイン実装の意義 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ディレクトリにコメント付きで収録されています。

タグ: PyTorch ONNX TensorRT CustomPlugin

7月24日 23:16 投稿