Есть ли способ, который может преобразовать формат данных файла .pb из NCHW в NHW C?

Question

У меня есть модель CNN, которая обучалась в Pytorch на основе формата данных N (партия) x C (канал) x H (высота) x W (ширина) . Я сохранил предварительно обученную модель как model.pth . После этого я преобразовал предварительно обученную модель из model.pth -> model.onnx , используя существующую функцию:

torch.onnx.export(model, dummy_input, "model.onnx")

И затем я преобразовал этот model.onnx -> model.pb с помощью модуля ниже:

import onnx
from onnx_tf.backend import prepare 

model_onnx = onnx.load('model.onnx')
tf_rep = prepare(model_onnx)
tf_rep.export_graph('model.pb')

Проблема заключается в следующем: я хочу использовать этот model.pb на устройстве с процессором, которому требуется NHW C формат данных. Тем не менее, моя модель основана на формате данных NCHW. Есть ли какой-нибудь метод, который может преобразовать формат данных этой модели .pb из NCHW в NHW C?

Answer 1

Можете ли вы просто транспонировать входные тензоры? например, input.transpose(1,2).transpose(2,3)?

>>> torch.randn( (3,3,3,3), names=['n','c','h','w']).transpose(1,2).transpose(2,3).names
('n', 'h', 'w', 'c')