31 lines
886 B
Python
31 lines
886 B
Python
import torch
|
|
import torchreid
|
|
|
|
archivo_salida = "osnet_x1_0_msmt17_batch1.onnx"
|
|
|
|
print("1. Descargando los pesos originales de la Bestia (OSNet x1_0)...")
|
|
# ⚡ Al poner pretrained=True, Python usa el enlace directo interno de la librería
|
|
# saltándose los bloqueos del navegador web.
|
|
model = torchreid.models.build_model(
|
|
name='osnet_x1_0',
|
|
num_classes=1000,
|
|
loss='softmax',
|
|
pretrained=True
|
|
)
|
|
model.eval()
|
|
|
|
print("2. Compilando el modelo ONNX blindado a Batch=1...")
|
|
dummy_input = torch.randn(1, 3, 256, 128)
|
|
|
|
torch.onnx.export(
|
|
model,
|
|
dummy_input,
|
|
archivo_salida,
|
|
export_params=True,
|
|
opset_version=18, # Mantenemos la versión 18 que ya nos funcionó perfecto
|
|
do_constant_folding=True,
|
|
input_names=['images'],
|
|
output_names=['features']
|
|
)
|
|
|
|
print(f"¡Éxito total! Tu modelo de grado comercial máximo está listo: {archivo_salida}") |