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