pytorch——保存‘类别名与类别数量’到权值文件中

前言

不知道大家有没有像我一样,每换一次不一样的模型,就要输入不同的num_classes和name_classes,反正我是很头疼诶,尤其是项目里面不止一个模型的时候,更新的时候看着就很头疼,然后就想着直接输入模型权值文件的path该多好,然后我就搞起来了。

在自己的类中加入想要加入数据信息

python 复制代码
class your_nets(nn.Module):
    def __init__(self, num_classes = 21,name_classes=None):
        super(your_nets, self).__init__()
        self.num_classes = num_classes
        self.name_classes = name_classes

训练过程之保存文件

复制代码
      
python 复制代码
model = your_nets(num_classes=num_classes, name_classes=name_classes)

save_dict = {
                'state_dict': model.state_dict(),
                'num_classes': model.num_classes,
                'name_classes': model.name_classes
            }

torch.save(save_dict, os.path.join(save_dir, "best_epoch_weights.pth"))

使用

python 复制代码
model = get_nets_class(model_path=model_path)


class get_nets_class(object):
    def __init__(self ,**kwargs):
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        load_dict  = torch.load(self.model_path, map_location=device)

        state_dict =load_dict['state_dict']
        num_classes = load_dict['num_classes']
        name_classes = load_dict['name_classes']

        if num_classes is not None and name_classes is not None:
            self.num_classes =num_classes
            self.name_classes = name_classes
            self.net = your_nets(num_classes=self.num_classes,name_classes=name_classe)
            self.net.load_state_dict(state_dict)
        else:
            self.net = your_nets(num_classes=self.num_classes, backbone=self.backbone)
            self.net.load_state_dict(load_dict)
        self.net = self.net.eval()
    
    def predict(self,image,name_classes,object_list):
        #你的预处理操作,没有就忽略
        image_data = preprocess(image)
        with torch.no_grad():
            # 推理
            pr = self.net(images)[0]
            # softmax 得出概率 pr.permute(1, 2, 0), dim=-1为我自己的操作,没有请忽略
            pr = F.softmax(pr.permute(1, 2, 0), dim=-1).cpu().numpy()
        #你的后处理操作,没有就忽略
        pr = postprocess(pr)
        #这一步与object_list有关 object_list是你想要模型去预测的内容
        # 例如你训练了识别cat、dog、pig、person的类别 那么你想只识别人,那么就object_list=['person'] 
        if object_list is not None:
            model_object_list = [name_classes.index(i) for i in object_list if i in name_classes]
            temp_list = [i for i in range(len(name_classes))]
            remove_list = [i for i in temp_list if i not in model_object_list]
            for i in remove_list:
                pr[pr==i] = 0
        retuen pr

我是觉得已经很详细了,大家要是不懂可以再问,我可以慢慢改进,每个人的写法都不一样 。

欢迎大家点赞加收藏哟~

相关推荐
聚客AI6 分钟前
搜索引擎vs向量数据库:LangChain混合检索架构实战解析
人工智能·pytorch·语言模型·自然语言处理·数据分析·gpt-3·文心一言
云畅新视界19 分钟前
从 CODING 停服到极狐 GitLab “接棒”,软件研发工具市场风云再起
人工智能·gitlab
chao_78920 分钟前
二分查找篇——寻找旋转排序数组中的最小值【LeetCode】
python·线性代数·算法·leetcode·矩阵
一ge科研小菜鸡24 分钟前
人工智能驱动下的可再生能源气象预测:构建绿色能源时代的新大脑
人工智能·能源
高压锅_122035 分钟前
Cursor+Coze+微信小程序实战: AI春联生成器
人工智能·微信小程序·notepad++
XiaoQiong.Zhang36 分钟前
数据分析框架和方法
人工智能
金玉满堂@bj37 分钟前
PyCharm 中 Python 解释器的添加选项及作用
ide·python·pycharm
程序员三藏41 分钟前
如何使用Pytest进行测试?
自动化测试·软件测试·python·测试工具·职场和发展·测试用例·pytest
TY-20251 小时前
三、神经网络——网络优化方法
人工智能·深度学习·神经网络