最近项目需要将之前训练好的pytorch 模型转成tensorflow版本,用keras 搭建了模型,想试着迁移权重,除了用onnx 直接转之外,更加精细的手动控制不同网络层的权重迁移。
下面以一个lstm 为例:
在 tf keras 下, 其中LSTM 和 Bidirectional 从keras导入即可:
由于是双向,所以一共包含 ih,hh,bias, ih_r, hh_r, bias_r 六个参数
而在 pytorch 的lstm 中 bias 的定义被拆成了两部分,
但维数相同,所以导出时,必须要加和才正好与tf.keras.layers 中的lstm bias相对应(这里一定注意!)
创建一个 空列表 将权重矩阵 从 torch.tensor 转换为 numpy, 并且要取一次转置(这里一定注意!)
layers对象都有一个 接口 set_weights(),传入提取的参数列表即可。
至此迁移 一个双向 lstm 层的 工作 就完成了。