import torch
from mmpose.models.necks import ChannelMapper
model = ChannelMapper(
in_channels=[256, 256],
kernel_size=1,
out_channels=384,
act_cfg=None,
norm_cfg=dict(type='BN'),
num_outs=3
)
# Create dummy input tensor with appropriate shape
inputs = (torch.randn(1, 256, 32, 32), torch.randn(1, 256, 16, 16))
dummy_input = (inputs,) # Wrap the inputs in an extra tuple
# Test the model with the dummy input
output = model.forward(dummy_input[0])
print(output[0].shape)
print(output[1].shape)
print(output[2].shape)
# Export to ONNX
torch.onnx.export(
model,
dummy_input, # Pass the wrapped dummy_input
"channel_mapper.onnx",
input_names=['input_0', 'input_1'],
output_names=['output_0', 'output_1', 'output_2'],
opset_version=11,
dynamic_axes={
'input_0': {2: 'height_0', 3: 'width_0'},
'input_1': {2: 'height_1', 3: 'width_1'},
'output_0': {2: 'out_height_0', 3: 'out_width_0'},
'output_1': {2: 'out_height_1', 3: 'out_width_1'},
'output_2': {2: 'out_height_2', 3: 'out_width_2'}
}
)
print("Model exported to channel_mapper.onnx")
result = self.forward(*input, **kwargs)
TypeError: ChannelMapper.forward() takes 2 positional arguments but 3 were given
이 에러가 떴다,, 추가 튜플로 래핑해서 model.forward에 단일 인수를 전달해서 예상 signature와 일치하게 됐음,,
'Programming > AI development' 카테고리의 다른 글
nvidia docker torch, trt runtime vs dev (0) | 2025.02.19 |
---|---|
raw image viewer (0) | 2025.01.11 |
Algorithm profiling (0) | 2024.09.06 |
머신비젼 사례 (1) | 2024.09.05 |
H100 vs RTX4090 (1) | 2024.08.26 |