File size: 2,617 Bytes
7382c66 20572f4 7382c66 20572f4 7382c66 20572f4 7382c66 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
import sys
import os
import torch
import json
from safetensors.torch import load_file
# Add src to path
sys.path.append(os.path.join(os.getcwd(), 'DA-2-repo/src'))
try:
from da2.model.spherevit import SphereViT
except ImportError as e:
print(f"Error importing SphereViT: {e}")
sys.exit(1)
# Load config
config_path = 'DA-2-repo/configs/infer.json'
with open(config_path, 'r') as f:
config = json.load(f)
# Adjust config for fixed size export
# Using 1092x546 (multiples of 14: 1092=78*14, 546=39*14)
# This is closer to the original config's ~600k pixels
H, W = 546, 1092
config['inference']['min_pixels'] = H * W
config['inference']['max_pixels'] = H * W
print(f"Initializing model with input size {W}x{H}...")
# Instantiate model
model = SphereViT(config)
# Load weights
print("Loading weights from model.safetensors...")
try:
weights = load_file('model.safetensors')
missing, unexpected = model.load_state_dict(weights, strict=False)
if missing:
print(f"Missing keys: {len(missing)}")
# print(missing[:5])
if unexpected:
print(f"Unexpected keys: {len(unexpected)}")
# print(unexpected[:5])
except Exception as e:
print(f"Error loading weights: {e}")
sys.exit(1)
print("Exporting model in FP32 (full precision)...")
model.eval()
# Dummy input (float32)
dummy_input = torch.randn(1, 3, H, W)
# Export
output_file = "onnx/model.onnx"
print(f"Exporting to {output_file}...")
try:
torch.onnx.export(
model,
dummy_input,
output_file,
opset_version=17,
input_names=["pixel_values"],
output_names=["predicted_depth"],
dynamic_axes={
"pixel_values": {0: "batch_size"},
"predicted_depth": {0: "batch_size"}
},
export_params=True,
do_constant_folding=True,
verbose=False
)
print(f"Successfully exported to {output_file}")
# Quantize the exported ONNX model
try:
from onnxruntime.quantization import quantize_dynamic, QuantType
quantized_output_file = "onnx/model_quantized.onnx"
print(f"Quantizing model to {quantized_output_file}...")
quantize_dynamic(
output_file,
quantized_output_file,
weight_type=QuantType.QInt8
)
print(f"Successfully quantized to {quantized_output_file}")
except Exception as qe:
print(f"Error during quantization: {qe}")
import traceback
traceback.print_exc()
except Exception as e:
print(f"Error exporting to ONNX: {e}")
import traceback
traceback.print_exc()
|