a tip for small memory computer failed run export_to_onnx.py when encoder = 'vitl'

my computer's memory is 16gb. When I run python export_to_onnx.py (encoder = 'vitl'), the command line output like this below,

username@username:~/depth-anything-tensorrt$ python export_to_onnx.py 
xFormers not available
xFormers not available
Total parameters: 335.32M
/home/username/depth-anything-tensorrt/torchhub/facebookresearch_dinov2_main/dinov2/layers/patch_embed.py:73: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  assert H % patch_H == 0, f"Input image height {H} is not a multiple of patch height {patch_H}"
/home/username/depth-anything-tensorrt/torchhub/facebookresearch_dinov2_main/dinov2/layers/patch_embed.py:74: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  assert W % patch_W == 0, f"Input image width {W} is not a multiple of patch width: {patch_W}"
/home/username/depth-anything-tensorrt/torchhub/facebookresearch_dinov2_main/vision_transformer.py:183: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if npatch == N and w == h:
/home/username/depth-anything-tensorrt/depth_anything/dpt.py:131: TracerWarning: Converting a tensor to a Python integer might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  out = F.interpolate(out, (int(patch_h * 14), int(patch_w * 14)), mode="bilinear", align_corners=True)
Killed

When I checked the resource usage of my computer while the code was running, I found that the memory was full. So, if your computer have a nvidia gpu, you can try this:

import os
import torch
import torch.onnx

from depth_anything.dpt import DPT_DINOv2
from depth_anything.util.transform import Resize, NormalizeImage, PrepareForNet

encoder = 'vitl'
#Download from https://huggingface.co/spaces/LiheYoung/Depth-Anything/tree/main/checkpoints
load_from = './checkpoints/depth_anything_vitl14.pth'
image_shape = (3, 518, 518)

# Initializing model
assert encoder in ['vits', 'vitb', 'vitl']
if encoder == 'vits':
    depth_anything = DPT_DINOv2(encoder='vits', features=64, out_channels=[48, 96, 192, 384], localhub='localhub')
elif encoder == 'vitb':
    depth_anything = DPT_DINOv2(encoder='vitb', features=128, out_channels=[96, 192, 384, 768], localhub='localhub')
else:
    depth_anything = DPT_DINOv2(encoder='vitl', features=256, out_channels=[256, 512, 1024, 1024], localhub='localhub')

total_params = sum(param.numel() for param in depth_anything.parameters())
print('Total parameters: {:.2f}M'.format(total_params / 1e6))

# Loading model weight
depth_anything = depth_anything.to('cuda')

depth_anything.load_state_dict(torch.load(load_from, map_location='cpu'), strict=True)

depth_anything.eval()

# Define dummy input data
dummy_input = torch.ones(image_shape).unsqueeze(0)

dummy_input = dummy_input.to('cuda')

# Provide an example input to the model, this is necessary for exporting to ONNX
example_output = depth_anything(dummy_input)

onnx_path = load_from.split('/')[-1].split('.pth')[0] + '.onnx'

# Export the PyTorch model to ONNX format

torch.onnx.export(depth_anything, dummy_input, onnx_path, opset_version=11, input_names=["input"], output_names=["output"], verbose=True)

print(f"Model exported to {onnx_path}")

compared to the source code, I added the following two lines:

depth_anything = depth_anything.to('cuda')
dummy_input = dummy_input.to('cuda')

spacewalk01 / depth-anything-tensorrt

a tip for small memory computer failed run export_to_onnx.py when encoder = 'vitl' #13