I’m very new to neural networks and Sentis. I made a python training script with help of ai that uses pretrained MobileNetV2 on my custom dataset. This dataset only has 3 categories which is black drawings on white background with 100 images in each. Training works and Ive tested model succesfully on the dataset itself and with small python drawing script. But when using it with unity sentis I’m getting same results no matter what images I use. If someone knows what could be the reason for this please let me know, I’ve spent almost 5 days already trying different things without any success.
Here’s my training script:
import os
import json
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, models
import numpy as np
DATASET_DIR = "dataset"
OUTPUT_DIR = "output"
IMAGE_SIZE = 224
BATCH_SIZE = 32
EPOCHS = 10
LEARNING_RATE = 0.0001
transform = transforms.Compose([
transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
def train():
os.makedirs(OUTPUT_DIR, exist_ok=True)
dataset = datasets.ImageFolder(DATASET_DIR, transform=transform)
loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
num_classes = len(dataset.classes)
print(f"Classes: {dataset.classes}")
print(f"Samples: {len(dataset)}\n")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}\n")
model = models.mobilenet_v2(weights='IMAGENET1K_V1')
model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
model = model.to(device)
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
criterion = nn.CrossEntropyLoss()
print("Training...")
for epoch in range(EPOCHS):
model.train()
correct = 0
total = 0
for images, labels in loader:
images, labels = images.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
correct += (outputs.argmax(1) == labels).sum().item()
total += labels.size(0)
accuracy = correct / total
print(f"Epoch {epoch+1}/{EPOCHS} - Accuracy: {accuracy:.3f}")
torch.save(model.state_dict(), "model.pth")
print("\n✓ Training complete")
export_to_onnx(model, dataset.classes)
def export_to_onnx(model, class_names):
print("\n" + "="*70)
print("EXPORTING TO ONNX")
print("="*70)
# Load and prepare model
model.load_state_dict(torch.load("model.pth", weights_only=True))
model.eval()
model.cpu()
# =========================================================================
# Wrap with normalization
# =========================================================================
print("Creating normalization wrapper...")
class NormWrapper(nn.Module):
def __init__(self, model):
super().__init__()
self.model = model
self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
def forward(self, x):
x = (x - self.mean) / self.std
return self.model(x)
wrapped = NormWrapper(model)
wrapped.eval()
print("✓ Model wrapped")
# =========================================================================
# Test PyTorch model
# =========================================================================
print("\nTesting PyTorch model...")
test_input = torch.ones(1, 3, IMAGE_SIZE, IMAGE_SIZE)
with torch.no_grad():
pytorch_out = wrapped(test_input)
print(f"PyTorch output: {pytorch_out[0].numpy()}")
if torch.isnan(pytorch_out).any() or torch.isinf(pytorch_out).any():
print("❌ ERROR: PyTorch output has NaN/Inf!")
return
if pytorch_out.abs().max() > 100:
print(f"❌ ERROR: PyTorch has extreme values ({pytorch_out.abs().max():.2f})!")
return
print("✓ PyTorch test passed")
# =========================================================================
# Export to ONNX
# =========================================================================
print("\nExporting to ONNX...")
dummy_input = torch.rand(1, 3, IMAGE_SIZE, IMAGE_SIZE)
torch.onnx.export(
wrapped,
dummy_input,
os.path.join(OUTPUT_DIR, "model.onnx"),
input_names=["input"],
output_names=["output"],
opset_version=25,
dynamo=True
)
print("✓ ONNX exported")
# =========================================================================
# Verify ONNX matches PyTorch
# =========================================================================
print("\nVerifying ONNX export...")
try:
import onnxruntime as ort
session = ort.InferenceSession(
os.path.join(OUTPUT_DIR, "model.onnx"),
providers=['CPUExecutionProvider']
)
# Test with same input
test_np = test_input.numpy()
onnx_out = session.run(None, {'input': test_np})[0]
pytorch_np = pytorch_out.numpy()
diff = np.abs(pytorch_np - onnx_out).max()
print(f"\nComparison:")
print(f" PyTorch: {pytorch_np[0]}")
print(f" ONNX: {onnx_out[0]}")
print(f" Max diff: {diff:.8f}")
if diff > 0.01:
print(f"\n❌ ERROR: Outputs don't match (diff={diff})!")
print("ONNX export is broken - DO NOT use in Unity!")
return
if np.abs(onnx_out).max() > 100:
print(f"\n❌ ERROR: ONNX has extreme values!")
return
print("\n✅ ONNX verification PASSED")
print("✅ Safe to import to Unity")
except ImportError:
print("\n⚠ WARNING: onnxruntime not installed")
print("Cannot verify export - install with: pip install onnxruntime")
except Exception as e:
print(f"\n❌ ERROR during verification: {e}")
return
# Save labels
with open(os.path.join(OUTPUT_DIR, "labels.json"), "w") as f:
json.dump(class_names, f)
print(f"\n{'='*70}")
print("EXPORT COMPLETE")
print(f"{'='*70}")
print(f"✓ Model: {OUTPUT_DIR}/model.onnx")
print(f"✓ Labels: {OUTPUT_DIR}/labels.json")
print(f"✓ Model expects 0-1 input (normalization built-in)")
print(f"{'='*70}")
if __name__ == "__main__":
train()
And unity script:
[SerializeField] private Texture2D _testImage1 = default!;
[SerializeField] private Texture2D _testImage2 = default!;
[SerializeField] private Texture2D _testImage3 = default!;
[SerializeField]
private ModelAsset _modelAsset = default!;
private Worker _worker = default!;
private Tensor<float> _inputTensor = default!;
private readonly string[] _classNames = ["earth", "fire", "water"];
private const int IMAGE_SIZE = 224;
private void Start()
{
var sourceModel = ModelLoader.Load(_modelAsset);
_worker = new Worker(sourceModel, BackendType.GPUCompute);
_inputTensor = new Tensor<float>(new TensorShape(1, 3, IMAGE_SIZE, IMAGE_SIZE));
TestImage(_testImage1, "Image 1");
TestImage(_testImage2, "Image 2");
TestImage(_testImage3, "Image 3");
}
private void TestImage(Texture2D texture, string name)
{
TextureConverter.ToTensor(texture, _inputTensor);
_worker.Schedule(_inputTensor);
var output = _worker.PeekOutput("output") as Tensor<float>;
var logits = output!.DownloadToArray();
output.Dispose();
Debug.Log($"=== Testing {name} === \nLogits: {string.Join(", ", logits.Select(v => v.ToString("F4")))} \nPredicted: {_classNames[Array.IndexOf(logits, logits.Max())]}");
}
This gives almost identical logits on 3 different images (imported to unity from dataset)
)