
java使用Deep Java Library(djl)搭配TorchScript搭建图片分类
java调用pytorch,深度学习应用,图像分类,deep java library,djl,Classification
·
一、前置要求
1.1、下载TorchScript类型的模型,注意这里是TorchScript类型,有些模型在说明中会说明是否为该格式的文件。可以从huggingface下载,在huggingface注意未区分PyTorch和TorchScript,在模型下方的标签都标记的为PyTorch,需要看具体的描述是否说该模型为TorchScript。
1.2、pom文件引入依赖,注意和引擎相关的包需要搭配引用,例如ai.djl.pytorch的native和jni包与engine版本要对上。pom.xml引入包如下:
<properties>
<maven.compiler.source>11</maven.compiler.source>
<maven.compiler.target>11</maven.compiler.target>
<djl.version>0.27.0</djl.version>
</properties>
<dependencies>
<!-- https://mvnrepository.com/artifact/ai.djl/api -->
<dependency>
<groupId>ai.djl</groupId>
<artifactId>api</artifactId>
<version>${djl.version}</version>
</dependency>
<!-- https://mvnrepository.com/artifact/ai.djl/model-zoo -->
<dependency>
<groupId>ai.djl</groupId>
<artifactId>model-zoo</artifactId>
<version>${djl.version}</version>
</dependency>
<!-- https://mvnrepository.com/artifact/ai.djl.pytorch/pytorch-engine -->
<dependency>
<groupId>ai.djl.pytorch</groupId>
<artifactId>pytorch-engine</artifactId>
<version>${djl.version}</version>
</dependency>
<dependency>
<groupId>ai.djl</groupId>
<artifactId>basicdataset</artifactId>
<version>${djl.version}</version>
</dependency>
<dependency>
<groupId>ai.djl.pytorch</groupId>
<artifactId>pytorch-engine</artifactId>
<version>${djl.version}</version>
</dependency>
<dependency>
<groupId>ai.djl.pytorch</groupId>
<artifactId>pytorch-jni</artifactId>
<version>2.1.1-0.27.0</version>
</dependency>
<dependency>
<groupId>ai.djl.pytorch</groupId>
<artifactId>pytorch-native-cpu</artifactId>
<classifier>win-x86_64</classifier>
<version>2.1.1</version>
</dependency>
<dependency>
<groupId>ai.djl</groupId>
<artifactId>djl-zero</artifactId>
<version>${djl.version}</version>
</dependency>
<dependency>
<groupId>org.apache.logging.log4j</groupId>
<artifactId>log4j-slf4j-impl</artifactId>
<version>2.21.0</version>
</dependency>
</dependencies>
二、java代码
将下载好的模型放到对应的位置,其中模型文件包含两个部分,一个是pt结尾的文件,当然结尾不一定是这个,可能是其他的,可以使用压缩文件打开这个模型文件看看是否包含了constants.pkl等文件,这个可以用作确认是否为TorchScript的一个标志。然后还需要一个synset.txt文件。
//这里也可以使用在线的模型
private static final URL MODEL_URL = NSFWUtil.class.getClassLoader().getResource("model/xxx.pt");
public static void main(String[] args) throws MalformedModelException, IOException, ModelNotFoundException, TranslateException {
getNSFW4JSON("image path");
}
/**
*
* @param imagePath 文件地址
* @throws ModelNotFoundException
* @throws MalformedModelException
* @throws IOException
* @throws TranslateException
* @return nsfw的json
*/
public static Classifications getNSFW4JSON(String imagePath) throws ModelNotFoundException, MalformedModelException, IOException, TranslateException {
Image img = ImageFactory.getInstance().fromFile(Paths.get(imagePath));
Translator<Image, Classifications> translator =
ImageClassificationTranslator.builder()
.addTransform(new Resize(224, 224))
.addTransform(new ToTensor())
.optApplySoftmax(true)
.build();
Criteria<Image, Classifications> criteria = Criteria.builder()
.setTypes(Image.class, Classifications.class)
.optModelUrls(MODEL_URL.toString())
.optTranslator(translator)
.optEngine("PyTorch") // Use PyTorch engine
.optProgress(new ProgressBar())
.build();
try (ZooModel<Image, Classifications> model = criteria.loadModel())
{
Predictor<Image, Classifications> predictor = model.newPredictor();
return predictor.predict(img);
}
}
/**
*
* @param in 输入流
* @throws ModelNotFoundException
* @throws MalformedModelException
* @throws IOException
* @throws TranslateException
* @return nsfw的json
*/
public static Classifications getNSFW4JSON(InputStream in) throws ModelNotFoundException, MalformedModelException, IOException, TranslateException {
Image img = BufferedImageFactory.getInstance().fromInputStream(in);
Translator<Image, Classifications> translator =
ImageClassificationTranslator.builder()
.addTransform(new Resize(224, 224))
.addTransform(new ToTensor())
.optApplySoftmax(true)
.build();
Criteria<Image, Classifications> criteria = Criteria.builder()
.setTypes(Image.class, Classifications.class)
.optModelUrls(MODEL_URL.toString())
.optTranslator(translator)
.optEngine("PyTorch") // Use PyTorch engine
.optProgress(new ProgressBar())
.build();
try (ZooModel<Image, Classifications> model = criteria.loadModel())
{
Predictor<Image, Classifications> predictor = model.newPredictor();
return predictor.predict(img);
}
}
三、总结
3.1、代码中的ImageClassificationTranslator在其他很多时候是自己在定义具体的方法实现,这里我们是图片分类,所以我们用的是官方提供的Translator。
3.2、就目前来说框架帮我们实现了很多的代码,能写的代码不是很多,难点在于如何找到能用的模型,目前很多模型还是PyTorch类型的,无法在JAVA或者C++环境调用。
3.3、可以试一下的模型nsfw,记住下synset.txt
更多推荐
所有评论(0)