TensorFlow Custom Allocator und Zugriff auf Daten von Tensor

Question

In TensorFlow können Sie benutzerdefinierte Zuweiser aus verschiedenen Gründen erstellen (ich mache es für neue Hardware). Aufgrund der Struktur des Geräts muss ich eine Struktur aus einigen Elementen als Datenzeiger verwenden, die der Zuordner als void* zurückgibt.

In den Kernen, die ich schreibe, bekomme ich Zugriff auf Tensors, aber ich brauche nicht die Pointer-Struktur, die ich geschrieben habe. Die Untersuchung der Klassen, so schien es, dass ich diese Struktur, indem Sie tensor_t.buf_->data()

Tensor::buf_

TensorBuffer::data()

Das Problem ist, bekommen konnte, dass ich nicht Code, der dies tut, finden kann, und ich bin besorgt, dass es unsicher ist, (sehr wahrscheinlich!) oder es gibt eine Standardmethode, dies zu tun.

Kann jemand bestätigen, ob dies eine gute/schlechte Idee ist? Und bieten Sie eine Alternative, wenn solche existiert?

Answer 1

Sie können auch Tensor::tensor_data() .data() verwenden, um auf den RAW-Zeiger zuzugreifen, ohne die seltsame Indirektion durch DMAHelper zu verwenden.

Answer 2

Vier Tage später ...

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

von GPUUtils

DMAHelper::base() ist ein Freund Klassenmethode, die die Fähigkeit gegeben wird, um den privaten Tensor::base() zu verwenden, um den Datenzeiger zu erhalten.

Die implementation zeigt, dass dies alles nur ein Wrapper um was ich nach einem anderen abstraction tun wollte. Ich vermute, es ist ein sicherer Ansatz, um den Zeiger zu bekommen und sollte stattdessen verwendet werden.