How To Fine-Tune Models for Better NSFW AI Detection?

Beginners
1

Hey everyone, I 'm new here. Hope you guys don’t mind answering basic questions.

I’ve been exploring NSFW AI detection lately, and it’s been a pretty fascinating rabbit hole. Tools like NSFWJS are great for quick setups, and the CLIP-based NSFW Detector is super impressive with how it uses embeddings to classify content.

Recently, I came across this site called soulfun.ai (which is all about creative AI stuff including ai generated photos and videos), and it got me thinking: how can I fine-tune these models for more niche or specific datasets?

I’ve been playing around with a basic CLIP setup, and here’s a quick snippet of what I’ve tried so far:

from transformers import CLIPProcessor, CLIPModel
import torch

# Load the pre-trained CLIP model
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# Set up inputs
inputs = processor(text=["NSFW", "SFW"], images=image, return_tensors="pt", padding=True)

# Forward pass
outputs = model(**inputs)
logits_per_image = outputs.logits_per_image  # Scores for image-text similarity
probs = logits_per_image.softmax(dim=1)  # Probabilities for each class

# Check if NSFW
is_nsfw = probs[0][0] > 0.5

Please let know, for those of you who’ve fine-tuned a CLIP-based model for NSFW (or even something similar):

  • What kind of datasets worked best for you?
  • Did you use any specific tricks during training to improve accuracy?
  • Any tips for keeping the model fast and lightweight during inference?

Would love to hear what’s worked for you! Thanks in advance for any advice. :blush:

1 Like
2

I found a document that describes some of the parameters used during the tuning process, even though it is a ViT model rather than a CLIP model.
The basic flow and libraries used are the same even when tuning a CLIP model. It’s just a different model class.

1 Like
4

Thanks, I’ll take a deep look!

1 Like
5

  • What kind of datasets worked best for you?
    I can suggest that datasets with a clear distinction between NSFW and SFW content, with a good balance of diverse examples, tend to work well. For instance, datasets like the one used by NSFWJS or any curated datasets from public sources that are labeled for NSFW content could be useful. If you have access to such datasets or can create one from scratch, that would be ideal.

  • Did you use any specific tricks during training to improve accuracy?
    Without running actual training, I can suggest some general techniques:

    • Data Augmentation: Slightly altering images (rotation, flipping, etc.) to increase dataset size and variability.

    • Transfer Learning: Using pre-trained models like CLIP as a starting point can leverage pre-existing knowledge, often improving accuracy.

    • Regularization: Techniques like dropout or L2 regularization to prevent overfitting.

    • Fine-tuning on Specific Classes: If your dataset has specific subcategories, fine-tuning on these can enhance model precision.

  • Any tips for keeping the model fast and lightweight during inference?

    • Model Pruning: Reduce the model size by removing less important weights or neurons.

    • Quantization: Convert the model to use lower precision (like INT8 instead of FP32) to reduce memory usage and increase speed.

    • Efficient Inference Frameworks: Use frameworks like ONNX Runtime or TensorRT for optimized inference.

    • Batch Processing: When possible, process multiple images or texts at once to leverage GPU efficiency.

2 Likes
6

The ViT model used a batch size of 16 and a learning rate of 5e-5. Do you think these parameters would be a good starting point for a CLIP model as well, or would adjustments be needed due to differences in model architecture? Anyway, thanks a lot!

1 Like
7

I don’t have much experience training models, so I don’t really know!:laughing:
However, since the image processing part of CLIP is ViT, I think it’s probably fine.
Well, I think that the optimal values are something that you have to try and find out, so I think it’s more reliable to adjust them while actually training.

1 Like
8

hahaha, I should stop being lazy and try it for myself, thanks. Well, training models and optimizing them is really like alchemy, I guess.

2 Likes
9

What other dataset do you use for such NSFW content detection apart from NSFWJS? Thanks

1 Like