Пользовательский распределитель TensorFlow и доступ к данным из тензора

В TensorFlow вы можете создавать собственные распределители по разным причинам (я делаю это для нового оборудования). Из-за структуры устройства мне нужно использовать структуру нескольких элементов в качестве моего указателя данных, который распределитель возвращает как void* .

В ядрах, которые я пишу, мне предоставляется доступ к тензорам, но мне нужно получить структуру указателя, которую я написал. Изучая classы, казалось, что я могу получить эту структуру, выполнив tensor_t.buf_->data()

Tensor::buf_

TensorBuffer::data()

Проблема в том, что я не могу найти код, который делает это, и я обеспокоен тем, что он небезопасен (весьма вероятно!), Или есть более стандартный способ сделать это.

Может ли кто-нибудь подтвердить, является ли это хорошей / плохой идеей? И предоставить альтернативу, если таковая существует?

Вы также можете использовать Tensor :: tensor_data () .data (), чтобы получить доступ к необработанному указателю, не используя странную косвенность через DMAHelper.

Четыре дня спустя …

 void* GetBase(const Tensor* src) { return const_cast(DMAHelper::base(src)); } 

от GPUUtils

DMAHelper::base() – метод classа друга, которому предоставляется возможность использовать закрытый Tensor::base() чтобы получить указатель данных.

Реализация показывает, что это всего лишь shell вокруг того, что я хотел сделать после еще одной абстракции . Я предполагаю, что это более безопасный подход к получению указателя и его следует использовать вместо этого.