pytorch & tensorflow 保存和加载模型

最新推荐文章于 2024-09-12 01:48:10 发布

原创最新推荐文章于 2024-09-12 01:48:10 发布 · 357 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #tensorflow #人工智能

代码专栏收录该内容

11 篇文章

订阅专栏

本文介绍了如何在PyTorch和TensorFlow中分别保存和加载模型结构与参数，包括使用`state_dict()`方法和`torch.save()`函数，以及`saver.save()`和`saver.restore()`方法来管理模型权重。

部署运行你感兴趣的模型镜像

1. Pytorch

1.1.1 save网络结构和参数：

注意最后一行为“self.state_dict()”

    def save(self,t):
        current_path = os.path.dirname(os.path.abspath(__file__))
        model_path = 'model/2E_model_' + t + '_'+self.name+'/'

        save_path = os.path.join(current_path,model_path)
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        save_file_path=os.path.join(save_path, 'model.pth')

        torch.save(self.state_dict(),save_file_path)

1.1.2 对应的加载模型参数：

注意对应“agent.load_state_dict(checkpoint)”

    def load(self,agent,model_path):
        model_pth = 'model.pth'
        model_path = os.path.join(model_path,model_pth)
        checkpoint = torch.load(model_path)
        agent.load_state_dict(checkpoint)
        agent.eval()

1.2.1 保存整个模型

注意为“torch.save(self.model,save_file_path)”

    def save(self,t):
        current_path = os.path.dirname(os.path.abspath(__file__))
        model_path = 'model/model_' + t + '_'+self.name+'/'

        save_path = os.path.join(current_path,model_path)
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        save_file_path=os.path.join(save_path, 'model.pth')

        torch.save(self.model,save_file_path)

1.2.2 加载整个模型

注意“self.model = torch.load(model_path)”

    def load(self,model_path):
        model_pth = 'model.pth'
        model_path = os.path.join(model_path,model_pth)
        self.model = torch.load(model_path)
        self.model.eval()

如果没对应上会报错：torch.nn.modules.module.ModuleAttributeError: object has no attribute 'copy'，参考此链接

2. Tensorflow

2.1 保存模型

    def save(self,time):
        current_path = os.path.dirname(os.path.abspath(__file__))
        model_path='model/model_'+time+'_'+self.name+'/weights_'+self.name
        save_path = os.path.join(current_path,model_path)
        if not os.path.exists(save_path):os.makedirs(save_path)
        self.saver.save(self.sess,save_path)

2.2 加载模型

    def load(self,model_path):
        meta_path = 'weights_'+self.name+'.meta'

        mata_path_dir = os.path.join(model_path,meta_path)

        self.saver = tf.compat.v1.train.import_meta_graph(mata_path_dir)
        a=model_path+'/'
        self.saver.restore(self.sess, tf.train.latest_checkpoint(a))

您可能感兴趣的与本文相关的镜像