27 lines
719 B
Python
27 lines
719 B
Python
|
import os
|
||
|
|
import torch
|
||
|
|
from torchvision.models.segmentation import deeplabv3_resnet50
|
||
|
|
|
||
|
|
resolution = (512, 512)
|
||
|
|
num_classes = 3
|
||
|
|
model_folder = "backup"
|
||
|
|
model_name = "ruasModel_best"
|
||
|
|
|
||
|
|
# Carrega seu modelo (usa num_classes igual ao treino)
|
||
|
|
model = deeplabv3_resnet50(weights=None, num_classes=num_classes)
|
||
|
|
model.load_state_dict(torch.load(os.path.join(model_folder, f"{model_name}.pth")))
|
||
|
|
model.eval()
|
||
|
|
|
||
|
|
# Dummy input com batch_size=1, channels=3, height=512, width=512
|
||
|
|
dummy_input = torch.randn(1, 3, resolution[0], resolution[1])
|
||
|
|
|
||
|
|
# Exporta
|
||
|
|
torch.onnx.export(
|
||
|
|
model,
|
||
|
|
dummy_input,
|
||
|
|
os.path.join(model_folder, f"{model_name}.onnx"),
|
||
|
|
input_names=["input"],
|
||
|
|
output_names=["output"],
|
||
|
|
opset_version=11
|
||
|
|
)
|