要在Java中調用PyTorch模型,可以使用PyTorch的Java API,也就是TorchScript。TorchScript是PyTorch的靜態圖編譯器,它允許將PyTorch模型編譯為一種可序列化和可導入的中間表示形式。然后可以在Java中加載并運行這個中間表示形式。
以下是一個簡單的示例代碼,展示了如何在Java中加載并調用一個PyTorch模型:
import org.pytorch.IValue;
import org.pytorch.Module;
import org.pytorch.Tensor;
public class PyTorchModelExample {
public static void main(String[] args) {
try {
// 加載PyTorch模型
Module module = Module.load("model.pt");
// 創建輸入Tensor
float[] inputData = {1.0f, 2.0f, 3.0f, 4.0f};
Tensor inputTensor = Tensor.fromBlob(inputData, new long[]{1, 4});
// 運行模型
IValue output = module.forward(IValue.from(inputTensor));
// 獲取輸出Tensor
Tensor outputTensor = output.toTensor();
// 打印輸出
float[] outputData = outputTensor.getDataAsFloatArray();
for (float value : outputData) {
System.out.println(value);
}
} catch (Exception e) {
e.printStackTrace();
}
}
}
在這個示例中,我們加載了一個名為model.pt
的PyTorch模型,并使用輸入數據{1.0, 2.0, 3.0, 4.0}
來運行模型。最后,我們獲取輸出Tensor并打印出來。
請注意,為了在Java中使用PyTorch的Java API,你需要在項目中添加PyTorch的Java庫依賴。你可以從PyTorch官方網站下載并添加到你的項目中。