Un exemple simple sur l'API de couche personnalisée (TensorRT 2.1)?

Question

J'utilise TensorRT 2.1 et je veux implémenter une couche personnalisée simple. (Le but est de lancer Single Shot Detector en utilisant TensorRT sur un système d'intégration.)

à la pratique, je voulais faire une couche Inc (juste en ajoutant 1,0 à une valeur de tenseur d'entrée et en gardant la même dimension).

Je implémente la classe Inc suivant le class Reshape : public Iplugin dans l'exemple sampleFasterRNN.cpp. J'ai tout gardé presque le même sauf getOutputDimensions() pour garder la même dimension. (Cela semble bien.)

Où devrais-je implémenter la partie "ajouter 1.0"? Je suppose que cela devrait être dans "enqueue()". Alors, j'ai essayé

int enqueue(int batchSize, const void*const *inputs, void** outputs, void*, cudaStream_t stream) override 
{ 
    # the below is from the Reshape class. seems to copy from input to output 
    CHECK(cudaMemcpyAsync(outputs[0], inputs[0], mCopySize * batchSize, cudaMemcpyDeviceToDevice, stream)); 
    # add 1.0 to first ten values 
    float* foutputs = (float*) outputs[0]; 
    int i; for (i = 0; i < 10; i++) foutputs[i] += 1.0; 
    return 0; 
} 

Cependant, cette partie entraîne erreur "erreur de segmentation".

Mes questions sont les suivantes:

1. Où et comment puis-je mettre en œuvre un calcul entre l'entrée et la sortie?
2. Quelqu'un peut-il fournir un exemple simple?

Answer 1

Référez-vous au fichier samples/samplePlugin/samplePlugin.cpp et regardons la classe FCPlugin. Votre calcul réel devrait aller dans la méthode enqueue. Vous devrez peut-être écrire un noyau CUDA qui fait l'incrément.