mmr/tkDNN
Archived
1
0
Fork 0
Code Issues 46 Pull Requests 12 Releases 4 Activity

LSTM cudnn test

Browse Source
This commit is contained in:
Francesco Gatti
2020-02-13 19:27:18 +01:00
parent 8a4d1cac17
commit a9c0db0bf6
10 changed files with 464 additions and 134 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -9,3 +9,4 @@ build/
*.tar.gz
*.weights
.idea/
*.hdf5
CMakeLists.txt
+3
View File
@@ -85,6 +85,9 @@ target_link_libraries(test_yolo3_berkeley tkDNN)
add_executable(test_yolo3_flir tests/yolo3_flir/yolo3_flir.cpp)
target_link_libraries(test_yolo3_flir tkDNN)
add_executable(test_imuodom tests/imuodom/imuodom.cpp)
target_link_libraries(test_imuodom tkDNN)
################################################################################
include/tkDNN/Layer.h
+76
View File
@@ -9,8 +9,10 @@
namespace tk { namespace dnn {
enum layerType_t {
LAYER_INPUT,
LAYER_DENSE,
LAYER_CONV2D,
LAYER_LSTM,
LAYER_ACTIVATION,
LAYER_FLATTEN,
LAYER_MULADD,
@@ -47,8 +49,10 @@ public:
std::string getLayerName() {
layerType_t type = getLayerType();
switch(type) {
case LAYER_INPUT: return "Input";
case LAYER_DENSE: return "Dense";
case LAYER_CONV2D: return "Conv2d";
case LAYER_LSTM: return "LSTM";
case LAYER_ACTIVATION: return "Activation";
case LAYER_FLATTEN: return "Flatten";
case LAYER_MULADD: return "MulAdd";
@@ -105,6 +109,27 @@ public:
};
/**
Input layer (it doesnt need weigths)
*/
class Input : public Layer {
public:
Input(Network *net, dataDim_t &dim, dnnType* srcData) : Layer(net) {
input_dim = dim;
output_dim = dim;
dstData = srcData;
}
virtual ~Input() {}
virtual layerType_t getLayerType() { return LAYER_INPUT; };
virtual dnnType* infer(dataDim_t &dim, dnnType* srcData) {
return dstData;
}
};
/**
Dense (full interconnection) layer
*/
@@ -148,6 +173,14 @@ protected:
/**
Convolutional 2D layer
WEIGHTS shape: OUTCH, INCH, KH, KW ...
BIAS shape: OUTCH
with BATCHNORM:
scales: OUTCH
means: OUTCH
variance: OUTCH
*/
class Conv2d : public LayerWgs {
@@ -172,6 +205,49 @@ protected:
size_t ws_sizeInBytes;
};
/**
Bidirectional LSTM layer
https://github.com/jiangnanhugo/seq2seq_cuda/blob/e4dbdcfa0517c972bfd4beea9f11a5233954093c/src/rnn.cpp
numlayers = 1 # hardcoded as 1
PARAMS (numlayers*2):
layer0:
( INCH, ? ) ???
( HIDDEN, ? ) ???
( HIDDEN * 8 ) ???
layer2:
( INCH, ? ) ???
( HIDDEN, ? ) ???
( HIDDEN * 8 ) ???
output shape: ( 2*HIDDEN, INH, INW )
*/
class LSTM : public Layer {
public:
LSTM(Network *net, int hiddensize, std::string fname_weights);
virtual ~LSTM();
virtual layerType_t getLayerType() { return LAYER_LSTM; };
virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
int kernelH, kernelW, strideH, strideW, paddingH, paddingW;
protected:
cudnnFilterDescriptor_t paramDesc;
cudnnTensorDescriptor_t hiddenStateTensorDesc, cellStateTensorDesc;
cudnnRNNDescriptor_t rnnDesc;
cudnnRNNDataDescriptor_t rnnDataDesc;
cudnnDropoutDescriptor_t dropDesc;
cudnnRNNAlgo_t algo;
dnnType *hiddenStateData, *cellStateData;
dnnType *paramsSpace;
void* workSpace;
size_t ws_sizeInBytes;
};
/**
Flatten layer
src/LSTM.cpp
+142
View File
@@ -0,0 +1,142 @@
#include <iostream>
#include "Layer.h"
namespace tk { namespace dnn {
LSTM::LSTM( Network *net, int hiddensize, std::string fname_weights) :
Layer(net) {
checkCUDNN( cudnnCreateFilterDescriptor(&paramDesc));
checkCUDNN( cudnnCreateRNNDescriptor(&rnnDesc) );
checkCUDNN( cudnnCreateRNNDataDescriptor(&rnnDataDesc) );
checkCUDNN( cudnnCreateDropoutDescriptor(&dropDesc));
int n = input_dim.n;
int c = input_dim.c;
int h = input_dim.h;
int w = input_dim.w;
checkCUDNN( cudnnSetTensor4dDescriptor(srcTensorDesc,
net->tensorFormat, net->dataType, n, 1, h, w) );
int numlayers = 1;
checkCUDNN( cudnnSetRNNDescriptor(net->cudnnHandle, rnnDesc, hiddensize, numlayers, dropDesc,
cudnnRNNInputMode_t::CUDNN_LINEAR_INPUT,
cudnnDirectionMode_t::CUDNN_BIDIRECTIONAL, cudnnRNNMode_t::CUDNN_LSTM,
cudnnRNNAlgo_t::CUDNN_RNN_ALGO_STANDARD, net->dataType) );
// find dimension of params
size_t params_size = 0;
checkCUDNN( cudnnGetRNNParamsSize(net->cudnnHandle, rnnDesc, srcTensorDesc, &params_size, net->dataType) );
std::cout<<"Params size bytes: "<<params_size<<", floats: "<<params_size/4<<"\n";
int dimW[3] = { int(params_size / sizeof(float)), 1, 1};
checkCUDNN(cudnnCreateFilterDescriptor(&paramDesc));
checkCUDNN(cudnnSetFilterNdDescriptor(paramDesc, net->dataType, net->tensorFormat, 3, dimW));
checkCuda( cudaMalloc(&paramsSpace, params_size) );
int numlinearlayers = 8;
for(int i=0; i<numlayers*2; i++) {
std::cout<<"layer: "<<i<<"\n";
for(int j=0; j<numlinearlayers; j++) {
// get weights pointer
cudnnFilterDescriptor_t linLayerMatDesc;
checkCUDNN(cudnnCreateFilterDescriptor(&linLayerMatDesc));
dnnType *linLayerMat;
checkCUDNN(cudnnGetRNNLinLayerMatrixParams(net->cudnnHandle, rnnDesc,
i, srcTensorDesc, paramDesc, paramsSpace,
j, linLayerMatDesc, (void **)&linLayerMat));
if(linLayerMat == nullptr) {
FatalError("LSTM No weights in hidden layer");
}
cudnnDataType_t dataType;
cudnnTensorFormat_t format;
int nbDims;
int filterDimA[3];
checkCUDNN(cudnnGetFilterNdDescriptor(linLayerMatDesc, 3, &dataType,
&format, &nbDims, filterDimA));
std::cout<<"Wgs Dims: "<<nbDims<<" ("<<filterDimA[0]<<", "<<filterDimA[1]<<", "<<filterDimA[2]<<")\n";
// here we should fill the params data into linLayerMat
checkCUDNN(cudnnDestroyFilterDescriptor(linLayerMatDesc));
// get bias pointer
cudnnFilterDescriptor_t linLayerBiasDesc;
checkCUDNN(cudnnCreateFilterDescriptor(&linLayerBiasDesc));
float *linLayerBias;
checkCUDNN(cudnnGetRNNLinLayerBiasParams(net->cudnnHandle, rnnDesc,
i, srcTensorDesc, paramDesc, paramsSpace,
j, linLayerBiasDesc, (void **)&linLayerBias));
if(linLayerMat == nullptr) {
FatalError("LSTM No bias in hidden layer");
}
checkCUDNN(cudnnGetFilterNdDescriptor(linLayerBiasDesc, 3, &dataType,
&format, &nbDims, filterDimA));
std::cout<<"bias Dims: "<<nbDims<<" ("<<filterDimA[0]<<", "<<filterDimA[1]<<", "<<filterDimA[2]<<")\n";
// here we should fill the params data into linLayerBiasDesc
checkCUDNN(cudnnDestroyFilterDescriptor(linLayerBiasDesc));
}
}
checkCUDNN( cudnnCreateTensorDescriptor(&hiddenStateTensorDesc));
checkCUDNN( cudnnSetTensor4dDescriptor(hiddenStateTensorDesc,
net->tensorFormat, net->dataType, 2*n, c, h, w) );
checkCuda( cudaMalloc(&hiddenStateData, 2*input_dim.tot()*sizeof(dnnType)) );
checkCUDNN( cudnnCreateTensorDescriptor(&cellStateTensorDesc));
checkCUDNN( cudnnSetTensor4dDescriptor(cellStateTensorDesc,
net->tensorFormat, net->dataType, 2*n, c, h, w) );
checkCuda( cudaMalloc(&cellStateData, 2*input_dim.tot()*sizeof(dnnType)) );
output_dim = input_dim;
output_dim.c = hiddensize*2;
checkCUDNN( cudnnSetTensor4dDescriptor(dstTensorDesc,
net->tensorFormat, net->dataType, output_dim.n, output_dim.c, output_dim.h, output_dim.w) );
//allocate data for infer result
checkCuda( cudaMalloc(&dstData, output_dim.tot()*sizeof(dnnType)) );
}
LSTM::~LSTM() {
checkCuda( cudaFree(dstData) );
}
dnnType* LSTM::infer(dataDim_t &dim, dnnType* srcData) {
checkCUDNN(cudnnRNNForwardInference(
net->cudnnHandle, rnnDesc, 1,
&srcTensorDesc, srcData,
hiddenStateTensorDesc, hiddenStateData,
cellStateTensorDesc, cellStateData,
paramDesc, paramsSpace,
&dstTensorDesc, dstData,
hiddenStateTensorDesc, hiddenStateData,
cellStateTensorDesc, cellStateData,
workSpace, ws_sizeInBytes
));
return dstData;
}
}}
src/utils.cpp
+2 -1
View File
@@ -39,7 +39,8 @@ void readBinaryFile(std::string fname, int size, dnnType** data_h, dnnType** dat
*data_h = new dnnType[size];
if (!dataFile.read ((char*) *data_h, size_b))
{
error_s << "Error reading file " << fname;
error_s << "Error reading file " << fname << " with n of float: "<<size;
error_s << " seek: "<<seek << " size: "<<size_b<<"\n";
FatalError(error_s.str());
}
tests/imuodom/imuodom.cpp
+70
View File
@@ -0,0 +1,70 @@
#include<iostream>
#include "tkdnn.h"
const char *i0_bin = "../tests/imuodom/layers/input0.bin";
const char *i1_bin = "../tests/imuodom/layers/input1.bin";
const char *i2_bin = "../tests/imuodom/layers/input2.bin";
const char *o0_bin = "../tests/imuodom/layers/output0.bin";
const char *o1_bin = "../tests/imuodom/layers/output1.bin";
const char *output_bin = "../tests/imuodom/layers/output.bin";
const char *c0_bin = "../tests/imuodom/layers/conv1d_7.bin";
const char *c1_bin = "../tests/imuodom/layers/conv1d_8.bin";
const char *c2_bin = "../tests/imuodom/layers/conv1d_9.bin";
const char *c3_bin = "../tests/imuodom/layers/conv1d_10.bin";
const char *c4_bin = "../tests/imuodom/layers/conv1d_11.bin";
const char *c5_bin = "../tests/imuodom/layers/conv1d_12.bin";
int main() {
// Network layout
tk::dnn::dataDim_t dim0(1, 4, 1, 100);
tk::dnn::dataDim_t dim1(1, 3, 1, 100);
tk::dnn::dataDim_t dim2(1, 3, 1, 100);
// Load input
dnnType *i0_d, *i1_d, *i2_d;
dnnType *i0_h, *i1_h, *i2_h;
readBinaryFile(i0_bin, dim0.tot(), &i0_h, &i0_d);
readBinaryFile(i1_bin, dim1.tot(), &i1_h, &i1_d);
readBinaryFile(i2_bin, dim2.tot(), &i2_h, &i2_d);
tk::dnn::Network net(dim0);
tk::dnn::Input x0 (&net, dim0, i0_d);
tk::dnn::Conv2d x0_0(&net, 128, 1, 11, 1, 1, 0, 0, c0_bin);
tk::dnn::Conv2d x0_1(&net, 128, 1, 11, 1, 1, 0, 0, c1_bin);
tk::dnn::Pooling x0_2(&net, 1, 3, 1, 3, tk::dnn::tkdnnPoolingMode_t::POOLING_MAX);
tk::dnn::Input x1 (&net, dim1, i1_d);
tk::dnn::Conv2d x1_0(&net, 128, 1, 11, 1, 1, 0, 0, c2_bin);
tk::dnn::Conv2d x1_1(&net, 128, 1, 11, 1, 1, 0, 0, c3_bin);
tk::dnn::Pooling x1_2(&net, 1, 3, 1, 3, tk::dnn::tkdnnPoolingMode_t::POOLING_MAX);
tk::dnn::Input x2 (&net, dim2, i2_d);
tk::dnn::Conv2d x2_0(&net, 128, 1, 11, 1, 1, 0, 0, c4_bin);
tk::dnn::Conv2d x2_1(&net, 128, 1, 11, 1, 1, 0, 0, c5_bin);
tk::dnn::Pooling x2_2(&net, 1, 3, 1, 3, tk::dnn::tkdnnPoolingMode_t::POOLING_MAX);
tk::dnn::Layer *concat_l[3] = { &x0_2, &x1_2, &x2_2 };
tk::dnn::Route concat (&net, concat_l, 3);
tk::dnn::LSTM lstm0(&net, 128, "ciao");
net.print();
dnnType *data;
tk::dnn::dataDim_t dim;
TIMER_START
// Inference
data = net.infer(dim, data); dim.print();
TIMER_STOP
// Print real test
std::cout<<"\n==== CHECK RESULT ====\n";
dnnType *out;
dnnType *out_h;
readBinaryFile(output_bin, dim.tot(), &out_h, &out);
checkResult(dim.tot(), data, out);
return 0;
}
tests/imuodom/infer.py
+74
View File
@@ -0,0 +1,74 @@
import keras
from keras.models import load_model
import keras.backend.tensorflow_backend as KTF
import numpy as np
import argparse
import tensorflow as tf
import os
import random
import struct
from keras.models import Sequential, Model
def bin_write(f, data):
data = data.flatten()
fmt = 'f'*len(data)
bin = struct.pack(fmt, *data)
f.write(bin)
if __name__ == '__main__':
print("DATA FORMAT: ", keras.backend.image_data_format())
print("Load model: ", "ferrariS1.hdf5")
model = load_model("ferrariS1.hdf5")
model.summary()
weights = model.get_weights()
x_angle = np.random.rand(1,100,4)
x_gyro = np.random.rand(1,100,3)
x_acc = np.random.rand(1,100,3)
[yhat_delta_p, yhat_delta_q] = model.predict([x_angle, x_gyro, x_acc], batch_size=1, verbose=1)
layer_name = 'bidirectional_3'
intermediate_layer_model = Model(inputs=model.input,
outputs=model.get_layer(layer_name).output)
intermediate_output = intermediate_layer_model.predict([x_angle, x_gyro, x_acc])
x_angle = np.array([x_angle])
x_gyro = np.array([x_gyro])
x_acc = np.array([x_acc])
intermediate_output = np.array([intermediate_output])
x_angle = x_angle.transpose(0, 3, 1, 2)
x_gyro = x_gyro.transpose(0, 3, 1, 2)
x_acc = x_acc.transpose(0, 3, 1, 2)
intermediate_output = intermediate_output.transpose(0, 3, 1, 2)
print("x0: ", np.shape(x_angle))
print("out: ",np.shape(intermediate_output))
x_angle = np.array(x_angle.flatten(), dtype=np.float32)
x_gyro = np.array(x_gyro.flatten(), dtype=np.float32)
x_acc = np.array(x_acc.flatten(), dtype=np.float32)
yhat_delta_p = np.array(yhat_delta_p.flatten(), dtype=np.float32)
yhat_delta_q = np.array(yhat_delta_q.flatten(), dtype=np.float32)
intermediate_output = np.array(intermediate_output.flatten(), dtype=np.float32)
f = open("layers/input0.bin", mode='wb')
bin_write(f, x_angle)
f = open("layers/input1.bin", mode='wb')
bin_write(f, x_gyro)
f = open("layers/input2.bin", mode='wb')
bin_write(f, x_acc)
f = open("layers/output0.bin", mode='wb')
bin_write(f, yhat_delta_p)
f = open("layers/output1.bin", mode='wb')
bin_write(f, yhat_delta_q)
f = open("layers/output.bin", mode='wb')
bin_write(f, intermediate_output)
tests/simple/test_model.py
+33 -27
View File
@@ -1,43 +1,49 @@
import keras
import numpy as np
from keras.models import Sequential
from keras.layers import Input, Dense, Activation, Flatten, Dropout, ELU, Reshape, Lambda
from keras.layers import Input, Dense, Activation, Flatten, Dropout, ELU, Reshape, Lambda, Conv1D
from keras.layers.convolutional import Convolution2D, Convolution3D
from keras.layers.pooling import MaxPooling2D, MaxPooling3D, AveragePooling3D
from keras.models import Sequential, Model
from keras.layers import Cropping2D
import keras.backend.tensorflow_backend as KTF
import struct
from keras.models import Sequential, Model
def dense_model():
model = Sequential()
def bin_write(f, data):
data = data.flatten()
fmt = 'f'*len(data)
bin = struct.pack(fmt, *data)
f.write(bin)
def create_model():
x1 = Input((6, 16), name='x1')
conv = Conv1D(4, 2)(x1)
model = Model([x1], [conv])
model.summary()
model.add(Reshape((10, 10, 1), input_shape=(10, 10)))
model.add(Convolution2D(2, (4, 4), subsample=(2, 2),
bias_initializer='random_uniform', activation="relu"))
model.add(Convolution2D(4, (2, 2), subsample=(1, 1),
bias_initializer='random_uniform', activation="relu"))
model.add(Flatten())
model.add(Dense(4, bias_initializer='random_uniform', activation="relu"))
sgd = keras.optimizers.Adam(lr=1e-4, decay=1e-8)
model.compile(optimizer=sgd, loss="mse")
return model
if __name__ == '__main__':
print "DATA FORMAT: ", keras.backend.image_data_format()
print ("DATA FORMAT: ", keras.backend.image_data_format())
model = dense_model()
model.save("net.h5")
model = create_model()
model.save("net.hdf5")
grid = np.random.rand(10,10)
X = grid[None,:,:]
i = np.array(grid.flatten(), dtype=np.float32)
print i
i.tofile("input.bin", format="f")
print "Input: ", X
x = np.random.rand(1,1,6,16)
r = model.predict( x[0], batch_size=1)
r = np.array([r])
x = x.transpose(0, 3, 1, 2)
r = r.transpose(0, 3, 1, 2)
print("in: ", np.shape(x))
print("out: ", np.shape(r))
x = np.array(x.flatten(), dtype=np.float32)
f = open("input.bin", mode='wb')
bin_write(f, x)
r = np.array(r.flatten(), dtype=np.float32)
f = open("output.bin", mode='wb')
bin_write(f, r)
r = model.predict( X, batch_size=1)
print np.shape(r)
print "Result: ", r
print "Result shape: ", np.shape(r)
r.tofile("output.bin", format="f")
tests/simple/test_simple.cpp
+3 -11
View File
@@ -2,23 +2,15 @@
#include "tkdnn.h"
const char *input_bin = "../tests/simple/input.bin";
const char *c0_bin = "../tests/simple/layers/c0.bin";
const char *c1_bin = "../tests/simple/layers/c1.bin";
const char *d2_bin = "../tests/simple/layers/d2.bin";
const char *c0_bin = "../tests/simple/layers/conv1d_1.bin";
const char *output_bin = "../tests/simple/output.bin";
int main() {
// Network layout
tk::dnn::dataDim_t dim(1, 1, 10, 10, 1);
tk::dnn::dataDim_t dim(1, 16, 1, 6);
tk::dnn::Network net(dim);
tk::dnn::Conv2d l0(&net, 2, 4, 4, 2, 2, 0, 0, c0_bin);
tk::dnn::Activation l1(&net, CUDNN_ACTIVATION_RELU);
tk::dnn::Conv2d l2(&net, 4, 2, 2, 1, 1, 0, 0, c1_bin);
tk::dnn::Activation l3(&net, CUDNN_ACTIVATION_RELU);
tk::dnn::Flatten l4(&net);
tk::dnn::Dense l5(&net, 4, d2_bin);
tk::dnn::Activation l6(&net, CUDNN_ACTIVATION_RELU);
tk::dnn::Conv2d l0(&net, 4, 1, 2, 1, 1, 0, 0, c0_bin);
// Load input
dnnType *data;
tests/weights_exporter.py
+51 -86
View File
@@ -5,96 +5,51 @@ import numpy as np
import argparse
import tensorflow as tf
import os
import msgpack
import lmdb
import random
import struct
from keras.models import Sequential, Model
def export_dense(name, weights, bias):
print "######## EXPORT", name, "LAYER ########"
print "Original weighs:"
print weights
print bias, "\n"
def bin_write(f, data):
data = data.flatten()
fmt = 'f'*len(data)
bin = struct.pack(fmt, *data)
f.write(bin)
#input, filters
I, C = np.shape(weights)
B = np.shape(bias)
print "w shape: ", I, C
print "b shape: ", B
def export_layer(name, weights, bias):
print ("######## EXPORT", name, "LAYER ########")
wgs = [ [ j[i] for j in weights ] for i in xrange(C) ]
wgs = np.array(wgs, dtype=np.float32)
print("wgs pretranpose: ", np.shape(weights))
# convert NHWC to NCHW
if(weights.ndim == 4):
weights = weights.transpose(3,2,0,1)
elif(weights.ndim == 3):
weights = weights.transpose(2,1,0)
else:
print("Ndim", weights.ndim)
raise("not implemented with dim" )
print "REPOSITIONED WEIGHTS:"
print wgs
print("weights: ", np.shape(weights))
print("bias: ", np.shape(bias))
weights = np.array(weights.flatten(), dtype=np.float32)
bias = np.array(bias, dtype=np.float32)
wgs.tofile(name + ".bin", format="f")
bias.tofile(name + ".bias.bin", format="f")
print "WEIGHTS saved\n"
print(len(weights) + len(bias))
def export_conv2d(name, weights, bias):
print "######## EXPORT", name, "LAYER ########"
print "Original weighs:"
print weights
print bias, "\n"
f = open(name + ".bin", mode='wb')
bin_write(f, weights)
bin_write(f, bias)
print ("WEIGHTS saved\n")
# height, width, input, filters
H, W, N, C = np.shape(weights)
B = np.shape(bias)
print "w shape: ", N, C, H, W
print "b shape: ", B
def export_bidir(name, weights):
print ("######## EXPORT", name, "LAYER ########")
wgs = weights.transpose()
wgs = wgs.transpose(0, 1, 3, 2)
print "Final shape:", np.shape(wgs)
wgs = np.array(wgs.flatten(), dtype=np.float32)
print "REPOSITIONED WEIGHTS:"
print wgs
bias = np.array(bias, dtype=np.float32)
wgs.tofile(name + ".bin", format="f")
bias.tofile(name + ".bias.bin", format="f")
print "WEIGHTS saved\n"
def export_conv3d(name, weights, bias):
print "######## EXPORT", name, "LAYER ########"
print "Original weighs:"
print weights
print bias, "\n"
print np.shape(weights)
# height, width, input, thickness, filters
H, W, T, N, C = np.shape(weights)
B = np.shape(bias)
print "w shape: ", T, C, H, W #thickness is number of images for cudnn
print "b shape: ", B
wgs = weights.transpose()
wgs = wgs.transpose(0, 1, 4, 3, 2)
print "Final shape:", np.shape(wgs)
wgs = np.array(wgs.flatten(), dtype=np.float32)
print "REPOSITIONED WEIGHTS:"
print wgs
bias = np.array(bias, dtype=np.float32)
wgs.tofile(name + ".bin", format="f")
bias.tofile(name + ".bias.bin", format="f")
print "WEIGHTS saved\n"
def get_session(gpu_fraction=0.5):
gpu_options = tf.GPUOptions(allow_growth=True)
#per_process_gpu_memory_fraction=gpu_fraction)
return tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
for w in weights:
print(np.shape(w))
#https://github.com/fchollet/keras/wiki/Converting-convolution-kernels-from-Theano-to-TensorFlow-and-vice-versa
if __name__ == '__main__':
KTF.set_session(get_session())
print("DATA FORMAT: ", keras.backend.image_data_format())
parser = argparse.ArgumentParser(description='KERAS WEIGHTS EXPORTER TO CUDNN')
parser.add_argument('model',type=str,
@@ -103,31 +58,41 @@ if __name__ == '__main__':
args = parser.parse_args()
print "DATA FORMAT: ", keras.backend.image_data_format()
print("DATA FORMAT: ", keras.backend.image_data_format())
print "Load model: ", args.model
print("Load model: ", args.model)
model = load_model(args.model)
model.summary()
weights = model.get_weights()
ws = np.shape(weights)
print "Weights shape:", ws
print("Weights shape:", ws)
if not os.path.exists(args.output):
os.makedirs(args.output)
num = 0
name_num = 0
for l in model.layers:
print("\n\nNAME: ", l.name)
print("input: ", l.input_shape, " output: ", l.output_shape)
wgs = l.get_weights()
print("wgs num: ", len(wgs))
name = l.name
if name.startswith("conv3d"):
export_conv3d(args.output + "/conv" + str(name_num), weights[num], weights[num+1])
export_layer(args.output + "/" + name, wgs[0], wgs[1])
elif name.startswith("conv2d"):
export_conv2d(args.output + "/conv" + str(name_num), weights[num], weights[num+1])
export_layer(args.output + "/" + name, wgs[0], wgs[1])
elif name.startswith("conv1d"):
export_layer(args.output + "/" + name, wgs[0], wgs[1])
elif name.startswith("dense"):
export_dense(args.output + "/dense" + str(name_num), weights[num], weights[num+1])
export_layer(args.output + "/" + name, wgs[0], wgs[1])
elif name.startswith("bidirectional"):
export_bidir(args.output + "/" + name, wgs)
else:
print "skip:", name, "has no weights"
print ("skip:", name, "has no weights")
continue
name_num += 1
num += 2