pytorch从头训练自己的图像分类模型并利用flask部署模型

本文推荐一个完整的pytorch图像分类代码，轻松上手，可以实现很好的分类性能。

本文所介绍的代码可以实现很好的图像分类的性能，可以在比赛中采用这个代码库，可以作为baseline使用，一般也能获得相对好一点的成绩，可以作为自己的baseline使用，强烈推荐这个代码仓库。

利用pytorch实现图像分类，其中包含的densenet，resnet，resnext，mobilenet，efficientnet等图像分类网络，可以根据需要再行利用torchvision扩展其他的分类算法。

近些天添加cnn与svm或者knn结合的pytorch实现，利用卷积神经网络提取特征，然后将特征保存为pkl文件，然后训练分类器，预测得出结果。

github代码地址：https://github.com/lxztju/pytorch_classification

如果有用欢迎star

实现功能

• 基础功能利用pytorch实现图像分类
• 包含带有warmup的cosine学习率优化
• 多模型融合预测，加权与投票融合
• 利用flask实现模型云端api部署
• 冻结相应层不训练（减轻过拟合程度）
• 实现多gpu训练，单gpu或者cpu预测
• label smooth的pytorch实现
• 添加cnn提取特征并使用svm，knn等分类器进行分类的pytorch代码

运行环境

• python3.7
• pytorch 1.1
• torchvision 0.3.0

代码仓库的使用

数据集形式

原始数据集存储形式为，同个类别的图像存储在同一个文件夹下，所有类别的图像存储在一个主文件夹datasets下。

利用preprocess.py将数据集格式进行转换（个人习惯这种数据集的方式）

转换后的数据集为，将训练集的路径与类别存储在train.txt文件中，测试机存储在val.txt中

模型介绍

仓库中模型densenet，mobilenet,resnext模型来自于torchvision

efficientnet来自于 https://github.com/lukemelas/EfficientNet-PyTorch

build_model.py 构建模型

训练

• 在cfg.py中修改合适的参数，并在train.py中选择合适的模型

python train.py

预测

在cfg.py中TRAINED_MODEL参数修改为指定的权重文件存储位置

python predict.py

• 多模型融合的使用方式类似于预测，使用multi-model_predict.py进行多模型融合

cnn + svm

代码 cnn_ml.py

flask云端部署

将训练存储好的权重文件，存储在flask_deployment文件夹中

然后修改server.py中路径运行即可
利用client.py进行调用