模型ckpt格式转换pytorch的pth格式

模型的checkpoint (ckpt) 文件格式通常是由TensorFlow创建的，而PyTorch的模型通常是保存为.pth格式。如果你有一个TF checkpoint文件想要转换成PyTorch可以使用的.pth文件，你可以按照以下步骤操作：

1. **安装必要的库**:
首先确保已安装`torch`, `tensorflow`, 和 `torch.utils.checkpoint` 或者`tf-nightly`（如果使用的是TensorFlow nightly版本）。

pip install torch tensorflow torch-utils-checkpoint

2. **加载TF checkpoint**:
使用TensorFlow将ckpt文件加载到一个`tf.train.Checkpoint`对象中。

import tensorflow as tf
tf_checkpoint = tf.train.Checkpoint(model=your_model)
status = tf_checkpoint.restore('path_to_your_ckpt')

3. **转换为PyTorch状态字典**:
现在你需要将`tf.train.Checkpoint`的内容转换为PyTorch的state_dict。

# 如果你的模型是一个PyTorch Module
state_dict = {k[len("model."):] : v.numpy() for k, v in tf_checkpoint.model.items()}

4. **保存为.pth文件**:
最后，你可以将这个state_dict保存为.pth文件。

torch.save(state_dict, 'path_to_save_your.pth')

5. **注意事项**:
- 注意名称映射，因为两个框架对层名可能有不同的约定。
- 如果有自定义层，它们的权重可能会丢失或需要手动处理。
- 对于变量形状不一致的情况，可能需要进行额外的数据调整。