1 command network inference
This commit is contained in:
@@ -5,16 +5,31 @@
|
||||
|
||||
namespace tkDNN {
|
||||
|
||||
struct dataDim_t;
|
||||
class Layer;
|
||||
const int MAX_LAYERS = 256;
|
||||
|
||||
class Network {
|
||||
|
||||
public:
|
||||
Network();
|
||||
virtual ~Network();
|
||||
|
||||
/**
|
||||
Do inferece for every added layer
|
||||
*/
|
||||
value_type* infer(dataDim_t &dim, value_type* data);
|
||||
|
||||
bool addLayer(Layer *l);
|
||||
|
||||
cudnnDataType_t dataType;
|
||||
cudnnTensorFormat_t tensorFormat;
|
||||
cudnnHandle_t cudnnHandle;
|
||||
cublasHandle_t cublasHandle;
|
||||
|
||||
private:
|
||||
Layer* layers[MAX_LAYERS]; //contains layers of the net
|
||||
int num_layers; //current number of layers
|
||||
};
|
||||
|
||||
}
|
||||
|
||||
@@ -12,6 +12,9 @@ Layer::Layer(Network *net, dataDim_t in_dim) {
|
||||
|
||||
checkCUDNN( cudnnCreateTensorDescriptor(&srcTensorDesc) );
|
||||
checkCUDNN( cudnnCreateTensorDescriptor(&dstTensorDesc) );
|
||||
|
||||
if(!net->addLayer(this))
|
||||
FatalError("Net reached max number of layers");
|
||||
}
|
||||
|
||||
Layer::~Layer() {
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
#include <iostream>
|
||||
|
||||
#include "Network.h"
|
||||
#include "Layer.h"
|
||||
|
||||
namespace tkDNN {
|
||||
|
||||
@@ -12,6 +13,8 @@ Network::Network() {
|
||||
|
||||
checkCUDNN( cudnnCreate(&cudnnHandle) );
|
||||
checkERROR( cublasCreate(&cublasHandle) );
|
||||
|
||||
num_layers = 0;
|
||||
}
|
||||
|
||||
Network::~Network() {
|
||||
@@ -20,4 +23,21 @@ Network::~Network() {
|
||||
checkERROR( cublasDestroy(cublasHandle) );
|
||||
}
|
||||
|
||||
value_type* Network::infer(dataDim_t &dim, value_type* data) {
|
||||
|
||||
//do infer for every layer
|
||||
for(int i=0; i<num_layers; i++)
|
||||
data = layers[i]->infer(dim, data);
|
||||
|
||||
return data;
|
||||
}
|
||||
|
||||
bool Network::addLayer(Layer *l) {
|
||||
if(num_layers == MAX_LAYERS)
|
||||
return false;
|
||||
|
||||
layers[num_layers++] = l;
|
||||
return true;
|
||||
}
|
||||
|
||||
}
|
||||
+5
-1
@@ -29,11 +29,15 @@ int main() {
|
||||
TIMER_START
|
||||
|
||||
// Inference
|
||||
data = net.infer(dim, data); dim.print();
|
||||
|
||||
/*
|
||||
//old Inference method
|
||||
data = c0.infer(dim, data); dim.print();
|
||||
data = a0.infer(dim, data); dim.print();
|
||||
data = c1.infer(dim, data); dim.print();
|
||||
data = a1.infer(dim, data); dim.print();
|
||||
|
||||
*/
|
||||
TIMER_STOP
|
||||
// Print result
|
||||
printDeviceVector(dim.tot(), data);
|
||||
|
||||
Reference in New Issue
Block a user