修正张量流检验点中张量的形状

时间:2018-01-07 14:20:27

标签: tensorflow tf-slim

我有一个张量流检查点,我可以在使用常规例程tf.train.Saver()saver.restore(session, 'my_checkpoint.ckpt')重新定义与其对应的图表后加载。

但是,现在,我想修改网络的第一层,以接受形状为[200, 200, 1]而不是[200, 200, 10]的输入。

为此,我想通过对[3, 3, 10, 32](3x3内核,10个输入通道,32个输出通道)到[3, 3, 1, 32]的第一层对应的张量形状进行修改。第三维。

我怎么能这样做?

2 个答案:

答案 0 :(得分:1)

我找到了一种方法,但不是那么简单。 给定一个检查点,我们可以将它转换为序列化的numpy数组(或者我们可能认为适合保存numpy数组字典的任何其他格式),如下所示:

checkpoint = {}
with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    saver.restore(sess, 'my_checkpoint.ckpt')
    for x in tf.global_variables():
        checkpoint[x.name] = x.eval()
    np.save('checkpoint.npy', checkpoint)

可能会有一些例外情况要处理,但让代码保持简单。

然后,我们可以在numpy数组上执行我们喜欢的任何操作:

checkpoint = np.load('checkpoint.npy')
checkpoint = ...
np.save('checkpoint.npy', checkpoint)

最后,我们可以在构建图表后手动加载权重:

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    checkpoint = np.load('checkpoint.npy').item()
    for key, data in checkpoint.iteritems():
        var_scope = ... # to be extracted from key
        var_name = ...  # 
        with tf.variable_scope(var_scope, reuse=True):
            var = tf.get_variable(var_name)
            sess.run(var.assign(data))

如果有一种更简单的方法,我全都耳朵!

答案 1 :(得分:0)

可以使用tensorflow::BundleReader读取源ckpt,使用tensorflow::BundleWriter重写。

tensorflow::BundleReader reader(Env::Default(), model_path_prefix);
std::vector<std::string> tensor_names;
reader.Seek("");
reader.Next();
for (; reader.Valid(); reader.Next()) {
    tensor_names.emplace_back(reader.key());
}
tensorflow::BundleWriter writer(Env::Default(), new_model_path_prefix);   
for (auto &tensor_name : tensor_names) {
        DataType dtype;
        TensorShape shape;        
        
        reader.LookupDtypeAndShape(tensor_name, &dtype, &shape);
        Tensor val(dtype, shape);
        Status bool_ret  = reader.Lookup(tensor_name, &val);
        std::cout << tensor_name << " " << val.DebugString() << std::endl;
        // modify dtype and shape. padding Tensor
        TensorSlice slice(new_shape.dims());
        writer.AddSlice(tensor_name, new_shape, slice, new_val);
    }
}
writer.Finish();