Przejdź do głównej zawartości

Model CV na Tritonie

1. Wytrenuj i Zapisz Swój Model PyTorch

Najpierw musisz mieć wytrenowany model PyTorch. Po wytrenowaniu modelu zapisz go za pomocą funkcji torch.save() z PyTorcha.

import torch

# Załóżmy, że `model` to Twój wytrenowany model PyTorch
model_path = "path/to/your/model.pth"
torch.save(model, model_path)

2. Konwertuj Model do Formatu ONNX

Serwer inferencji Triton obsługuje wiele formatów modeli, w tym ONNX, TensorFlow i TensorRT. ONNX (Open Neural Network Exchange) jest popularnym wyborem ze względu na swoją kompatybilność z wieloma platformami i frameworkami.

Aby przekonwertować model PyTorch na ONNX:

import torch.onnx

# Załaduj wytrenowany model
model = torch.load('path/to/your/model.pt')
model.eval()

# Stwórz przykładowe dane wejściowe w kształcie oczekiwanym przez model
dummy_input = torch.randn(1, 3, 224, 224) # przykładowy kształt wejścia dla modelu obrazowego

# Eksportuj model
onnx_model_path = "path/to/exported/model.onnx"
torch.onnx.export(model, dummy_input, onnx_model_path)

3. Stwórz Plik Konfiguracyjny dla Tritona

Triton wymaga pliku config.pbtxt, aby zrozumieć, jak obsłużyć Twój model. Ten plik zawiera informacje o nazwie modelu, platformie (ONNX, TensorFlow itp.), tensorach wejściowych i wyjściowych oraz innych ustawieniach.

Oto przykładowy config.pbtxt dla modelu ONNX:

name: "model_name"
platform: "onnxruntime_onnx"
input [
{
name: "images0" # Nazwa warstwy wejściowej w Twoim modelu
data_type: TYPE_FP32
dims: [3, 224, 224] # Zastąp wymiarami wejściowymi Twojego modelu
}
]
output [
{
name: "output0" # Nazwa warstwy wyjściowej w Twoim modelu
data_type: TYPE_FP32
dims: [6, 17661] # Zastąp wymiarami wyjściowymi Twojego modelu
}
]

Możesz użyć netron.app, aby zobaczyć szczegóły Twojego modelu (wymiary wejścia, wymiary wyjścia itp.).

4. Wdróż Model na Tritonie

Umieść plik ONNX modelu oraz plik config.pbtxt w katalogu wewnątrz katalogu modeli Tritona. Struktura katalogów powinna wyglądać tak:

└──/models-repository
└──/model_name
├──config.pbtxt
└──/1
└──model_name.onnx

5. Uruchom Tritona

Aby uruchomić swój serwer inferencji Triton, proszę odnieść się do Uruchamianie