提高CNN层的theano功能的速度

我使用Keras训练了一个模型，并将tf作为后端：

activation = 'relu'
initializer = 'he_normal'

n_hidden = [256,128]
n_filters = [32]
input_shape = (batch_size,7213,1)


model = Sequential([

    InputLayer(batch_input_shape=input_shape),

    Convolution1D(nb_filter=n_filters[0], filter_length=8, activation=activation, border_mode='same', init=initializer, input_shape=input_shape),

    MaxPooling1D(pool_length=4),

    Flatten(),

    Dense(output_dim=n_hidden[0], activation=activation, init=initializer),

    Dense(output_dim=n_hidden[1], activation=activation, init=initializer),

    Dense(output_dim=3, input_dim=n_hidden[1], activation='linear'),
])

我需要构建一个使用我保存的权重复制model.predict（）的theano函数，以便返回w.r.t的输出的雅可比矩阵 . 输入 .

以下numpy代码提供与model.predict（）相同的结果：

pool_length=4

x_ = test_X_data.reshape(batch_size,7213)
weights_ = model.get_weights() #obtain model weights

#reshape CNN weights and bias weights
weights_[0] = np.reshape(weights_[0], (weights_[0].shape[0], weights_[0].shape[3]))
weights_[1] = np.reshape(weights_[1], (1, weights_[1].shape[0]))
weights_[3] = np.reshape(weights_[3], (1, weights_[3].shape[0]))
weights_[5] = np.reshape(weights_[5], (1, weights_[5].shape[0]))
weights_[7] = np.reshape(weights_[7], (1, weights_[7].shape[0]))

# pad left and right sides of input 
x_padded = np.pad(x_, ((0, 0), (3, 4)), mode='constant')

# compute Conv1d layer with bias weights
prediction = np.zeros((x_.shape[0],x_.shape[1],weights_[0].shape[1]))
for i in range(x_.shape[1]):
    prediction[:,i] = np.dot(x_padded[:,i:i+8],weights_[0])+weights_[1]
prediction = prediction

# RELU activation
prediction[prediction<0]=0

# Max pooling layer
pred_temp = np.zeros((prediction.shape[0],prediction.shape[1]/pool_length,prediction.shape[2]))
for i in range(prediction.shape[2]):
    for j in range(prediction.shape[1]/pool_length):
        pred_temp[:,j,i] = np.max(prediction[:,j*4:(j+1)*4,i],axis=1)

prediction = pred_temp.reshape(pred_temp.shape[0],pred_temp.shape[1]*pred_temp.shape[2])

# Dense layers
weights=np.vstack([weights_[2],weights_[3]])
prediction=np.hstack([prediction,np.zeros(prediction.shape[0]).reshape(prediction.shape[0],1)])
prediction[:,-1]=1
prediction=np.dot(prediction,weights)
prediction[prediction<0]=0

weights=np.vstack([weights_[4],weights_[5]])
prediction=np.hstack([prediction,np.zeros(prediction.shape[0]).reshape(prediction.shape[0],1)])
prediction[:,-1]=1
prediction=np.dot(prediction,weights)
prediction[prediction<0]=0


weights=np.vstack([weights_[6],weights_[7]])
prediction=np.hstack([prediction,np.zeros(prediction.shape[0]).reshape(prediction.shape[0],1)])
prediction[:,-1]=1
prediction= np.dot(prediction,weights)

下面是我尝试将其转换为计算Jacobian的theano函数，但我相信forloops使编译速度太慢 .

theano.config.optimizer='fast_compile'
theano.config.exception_verbosity='high'

weights_in_model = theano.typed_list.TypedListType(theano.tensor.dmatrix)()
x = T.matrix('x')

def pred_jac(x,weights_in_model): 

    pool_length=4

    x = T.concatenate((T.zeros((T.shape(x)[0], 3)),x,T.zeros((T.shape(x)[0], 4))))

    # Apply Convolution weights to input
    prediction = []
    for i in range(7213):
        prediction.append(T.dot(x[:,i:i+8],weights_in_model[0])+weights_in_model[1])
    prediction = T.as_tensor_variable(prediction)
    prediction = T.clip(prediction, 0, 9999.) # RELU activation

    prediction.dimshuffle(1,0,2) # Reformat to proper dimension order

    # Maxpooling layer
    pred_temp = []
    for i in range(32):
        pred_temp_b = []
        for j in range(1803):
            pred_temp_b.append(T.max(prediction[:,j*pool_length:(j+1)*pool_length,i],axis=1))
        pred_temp.append(T.as_tensor_variable(pred_temp_b))

    pred_temp = T.as_tensor_variable(pred_temp)

    pred_temp.dimshuffle(1,2,0)

    # Dense layers

    prediction = T.reshape(pred_temp,(T.shape(pred_temp)[0],T.shape(pred_temp)[1]*T.shape(pred_temp)[2]))


    weights = T.concatenate((weights_in_model[2],weights_in_model[3]), axis=0)
    prediction = T.concatenate((prediction, T.ones((T.shape(prediction)[0], 1))), axis=1)
    prediction = T.dot(prediction, weights)
    prediction = T.clip(prediction, 0, 9999.)

    weights = T.concatenate((weights_in_model[4],weights_in_model[5]), axis=0)
    prediction = T.concatenate((prediction, T.ones((T.shape(prediction)[0], 1))), axis=1)
    prediction = denormalize(T.dot(prediction, weights))
    prediction = T.flatten(prediction)
    return prediction

# Jacobian returns the first order partial derivatives of outputs w.r.t inputs
jac = theano.gradient.jacobian(pred_jac(x,weights_in_model),wrt=x)
compute_jac = theano.function([x,weights_in_model],[jac],allow_input_downcast=True)

有关如何改进此功能和/或加快函数编译和计算时间的任何建议？