Falconsai/nsfw_image_detection · How to convert it to TorchScript?

Feb 21

https://djl.ai/docs/pytorch/how_to_convert_your_model_to_torchscript.html

It seems i need some weights from training here:

example = torch.rand(1, 3, 224, 224)

DenisNovac

Feb 24

I found how to do this:

from PIL import Image
from transformers import AutoTokenizer

model_name = "Falconsai/nsfw_image_detection"

model = AutoModelForImageClassification.from_pretrained(model_name, torchscript=True, return_dict=False)

processor = AutoImageProcessor.from_pretrained(model_name)

# example nsfw picture to trigger the model, find it yourself
image = Image.open("images/hentai.jpg")
image_inputs = processor(images=image, return_tensors="pt")

config = {'forward': [image_inputs['pixel_values']]}
converted = torch.jit.trace_module(model,  config)

torch.jit.save(converted, "converted.pt")

synset.txt:

normal
nsfw

DenisNovac changed discussion status to closed Feb 24