Tensorflow预训练模型转PyTorch

深度学习领域是计算机科学中变化最快的领域之一。大约 5 年前,当我开始研究这个主题时,TensorFlow 被认为是主导框架。如今,大多数研究人员已经转向 PyTorch。

NSDT工具推荐Three.js AI纹理开发包 - YOLO合成数据生成器 - GLTF/GLB在线编辑 - 3D模型格式在线转换 - 可编程3D场景编辑器 - REVIT导出3D模型插件 - 3D模型语义搜索引擎 - Three.js虚拟轴心开发包 - 3D模型在线减面 - STL模型在线切割

虽然这种快节奏令人兴奋,但也带来了很多挑战。最近,我面临着继续完成 2018 年开展的一个项目的任务。一位同事在大量临床数据集上训练了一个分割模型,并报告了出色的性能。

今天,我们的目标是在称为迁移学习的过程中,将该训练好的模型用于类似的目标。这里的直觉是,与其从头开始,不如至少部分使用预训练权重来实例化新模型的权重,这将提供一个更好的起点。

1、收集 Tensorflow 1.x 权重

这听起来比实际容易。在 tensorflow 1.x 中,模型保存在四个单独的文件中 - 其中没有一个可以直接转换为 pytorch 的 state_dict。为了解决这个问题,我们必须手动创建一个字典并从 tensorflow 后端检索权重。

为了实现这一点,你需要了解 tensorflow 实现的命名方案。每个操作都可以在创建时分配一个名称。这个名称在稍后转换为 pytorch 时很重要。

import tensorflow as tf  # tensorflow 1.x
import pickle

'''
<base_folder>
├───checkpoint
├───<model_name>.meta
├───<model_name>.data-00000-of-00001
└───<model_name>.index
'''

# First let's load meta graph and restore weights
sess = tf.Session()
saver = tf.train.import_meta_graph(r'<base_folder>\<model_name>.meta')
saver.restore(sess, tf.train.latest_checkpoint(r'<base_folder>'))

# get all trainable weights and save them in a dictionary
vars = sess.graph.get_collection('trainable_variables')
weights = {}
for v in vars:
    weights[v.name] = sess.run(v)  # retrieve the value from the tf backend

with open('weights.pickle', 'wb') as handle:
    pickle.dump(weights, handle, protocol=pickle.HIGHEST_PROTOCOL)

2、重建模型

遗憾的是,没有直接的方法将 TensorFlow 模型转换为 PyTorch。但是,尽管语法略有不同,但大多数层都存在于这两个框架中。例如,在 tf1 中,卷积层可以包含激活函数,而在 PyTorch 中,该函数需要按顺序添加。

此示例展示了 tf1 和 PyTorch 实现中流行的 UNet 架构的 upconv 块。

# >>> tf1 implementation (without encapsulating class)
import tensorflow as tf

def upconvcat(self, x1, x2, n_filter, name):
    x1 = tf.keras.layers.UpSampling2D((2, 2))(x1)
    x1 = tf.layers.conv2d(x1, filters=n_filter, kernel_size=(3, 3), padding='same', name="upsample_{}".format(name))
    return tf.concat([x1, x2], axis=-1, name="concat_{}".format(name))  # NHWC format

# >>> pytorch implementation
import torch

class UpConvCat(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.up = torch.nn.Upsample(scale_factor=2)
        self.conv = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        return torch.cat([x1, x2], dim=1)  # NCHW format

3、NHWC 与 NCHW

tensorflow 和 pytorch 之间的最后一个重要区别是关于轴的约定。

在旧版 tensorflow 中, data_format 属性可以指定为 channels_lastchannels_first ,而前者是默认选项。然而,在 pytorch 中,只能使用 channels first。通常,这些格式表示为 NHWC 和 NCHW,分别表示批处理大小 (N)、高度 (H)、宽度 (W) 和通道 (C)。

np.transpose(kernel, (3, 2, 0, 1))

如果使用默认的 channels_lastoption 训练 tensorflow 中的预训练模型,则需要对内核轴进行置换才能与 torch 一起使用。为了弥补这一点,需要像这样调整 2d-conv 层权重。

4、初始化 pytorch 模型

将权重转换为正确的格式后,我们可以将它们加载到 pytorch 模型中。为此,我们随机实例化一个模型并遍历命名参数列表。然后我们使用来自 tensorflow 的权重就地修改参数。

    # set new weights from loaded tf values
    with torch.no_grad():
        for (name, param), (tf_name, tf_param) in zip(m.named_parameters(), tf_weights.items()):
            # convert NHWC to NCHW format and copy to change memory layout
            tf_param = np.transpose(tf_param, (3, 2, 0, 1)).copy() if len(tf_param.shape) == 4 else tf_param
            assert tf_param.shape == param.detach().numpy().shape, name

            # https://discuss.pytorch.org/t/how-to-assign-an-arbitrary-tensor-to-models-parameter/44082/3
            param.copy_(torch.tensor(tf_param, requires_grad=True, dtype=param.dtype))

5、结束语

按照这些步骤,可以提取在 tensorflow 1.x 中训练的模型并将其转换为 pytorch 模型。我希望这对与我处境相似的人有所帮助。


原文链接:TF预训练模型转PyTorch - BimAnt

相关推荐
Elastic 中国社区官方博客1 小时前
使用 Elastic AI Assistant for Search 和 Azure OpenAI 实现从 0 到 60 的转变
大数据·人工智能·elasticsearch·microsoft·搜索引擎·ai·azure
江_小_白2 小时前
自动驾驶之激光雷达
人工智能·机器学习·自动驾驶
yusaisai大鱼3 小时前
TensorFlow如何调用GPU?
人工智能·tensorflow
weixin_466202783 小时前
第31周:天气识别(Tensorflow实战第三周)
分类·数据挖掘·tensorflow
珠海新立电子科技有限公司6 小时前
FPC柔性线路板与智能生活的融合
人工智能·生活·制造
IT古董6 小时前
【机器学习】机器学习中用到的高等数学知识-8. 图论 (Graph Theory)
人工智能·机器学习·图论
曼城周杰伦6 小时前
自然语言处理:第六十三章 阿里Qwen2 & 2.5系列
人工智能·阿里云·语言模型·自然语言处理·chatgpt·nlp·gpt-3
余炜yw7 小时前
【LSTM实战】跨越千年,赋诗成文:用LSTM重现唐诗的韵律与情感
人工智能·rnn·深度学习
莫叫石榴姐7 小时前
数据科学与SQL:组距分组分析 | 区间分布问题
大数据·人工智能·sql·深度学习·算法·机器学习·数据挖掘
如若1237 小时前
利用 `OpenCV` 和 `Matplotlib` 库进行图像读取、颜色空间转换、掩膜创建、颜色替换
人工智能·opencv·matplotlib