classification

Pytorch 用于训练自己的图像分类模型

环境要求

pytorch opencv

conda虚拟环境一键导入：

conda env create -f torch.yaml

理论地址：

csdn 博客地址： MLP-Mixer：https://no-coding.blog.csdn.net/article/details/121682740

Conv-Mixer：https://blog.csdn.net/qq_38676487/article/details/120705254

ConvNeXt：https://blog.csdn.net/qq_38676487/article/details/123298605

How2Train

1.数据集:

─dataset
    ├─train
    │	└─cats
    │		└─xxjj.jpg
    │	└─dogs
    │	 	└─xxx.jpg
    ├─test
    │	└─cats
    │	└─dogs

classes.txt:

cats
dogs

在txt_annotation.py 中 calsses 与上述文件classes.txt 分类顺序一致，运行txt_annotation.py 生成 cls_train.txt, cls_text.txt

classes = ["cats", "dogs"]
sets    = ["train", "test"]

2.训练

​ 在config 中配置训练参数：

Cuda             = True  #是否使用GPU 没有为Flase

input_shape      = [48,48]  # 输入图片大小
is_grayscale     = False #灰度图训练为True 非灰度图为False
batch_size      = 2 # 自己可以更改
lr              = 1e-3         

classes_path    = 'classes.txt'


num_workers     = 0  # 是否开启多进程


train_annotation_path     = 'cls_train.txt'  # 训练集用于训练

val_annotation_path     = 'cls_test.txt'    # 测试集合用于每个epoch 训练完测试。保证训练属于充分




resume          =''  # 加载训练权重路径

log_dir         = 'logs' # 日志路径 tensorboard 保存

#------------------------------------------#
#   FocalLoss ：处理样本不均衡
#   alpha   
#   gamma >0 当 gamma=0 时就是交叉熵损失函数
#   论文中gamma = [0,0.5,1,2,5]
#   一般而言当γ增加的时候，a需要减小一点
#   reduction ： 就平均：'mean' 求和 'sum'
#   还未ti
#------------------------------------------#
#Focal_loss      = True  # True Focal loss 处理原本不均衡 False  使用 CrossEntropyLoss() # 还未使用成功

#label_smoothing 防止过拟合
label_smoothing =  True #

smoothing_value = 0.1  #[0，1] 之间



#学习率变化策略
scheduler   = 'cos' #[None,reduce,cos] None保持不变 reduce  按epoch 来减少 cos 余弦下降算法

在 trian.py 中

#---------------------------------------------------#
#  定义模型，可在nets 导入自己的模型去训练，
#  目前支持MLP-Mixer Conv-Mixer ConvNeXt系列模型
#  只有ConvNeXt 支持pretrain 官方提供的权重
#---------------------------------------------------#
 model = ConvMixer_768_32(n_classes=num_classes)

日志查看

由于每次启动训练时，会在logs 文件下按照时间创建一个日志文件。如：

tensorboard --logdir=logs\loss_2022_03_06_12_11_30

How2Eval

在eval.py 中:

if __name__ == "__main__":

    # 读取测试集路劲和标签
    with open("./cls_test.txt","r") as f: 
        lines = f.readlines()
    #---------------------------------------------------#
    #   权重和模型
    #   注意：训练时设置的模型需要和权重匹配，
    #   也就是训练的啥模型使用啥权重
    #---------------------------------------------------#
    model_path = '' #训练好的权重路径
    model = ConvMixer_768_32(n_classes=2) # 自己训练好的模型

    mode = load_dict(model_path,model) # 加载权重
    eval = eval_top(anno_lines=lines,model=model)
    #---------------------------------------------------#
    #   top1 预测概率最好高的值与真实标签一致 √
    #   top5 预测概率前五个值由一个与真实标签一致 √
    #---------------------------------------------------#
    print('start eval.....')
    top1 = eval.eval_top1()
    
    top5 = eval.eval_top5()
    print('top1:%.3f,top5:%3.f'%(top1,top5))
    print('Eval Finished')

How2Predict

predict.py 中，设置好模型和权重，控制台输入图片路径。

#加载模型
model_path = 'logs\ep050-loss0.414-val_loss0.376.pth'
model = ConvMixer_768_32(n_classes=2)
model = load_dict(model_path,model)
eval = eval_top(anno_lines=None,model=model)

while True:
    img = input('Input image filename:')
    try:
        image = Image.open(img)
    except:
        print('Open Error! Try again!')
        continue
    else:
        class_name = eval.detect_img(image,mode='predict')
        print(class_name)

控制台：

Loading weights into state dict...
Input image filename:d:\Classification\torch\datasets\test\cats\cat.4006.jpg 

训练技巧和练丹

• Focl_loss(样本不均衡策略)

• label_smoothing (训练样本偏少时，防止过拟合策略)

• 学习率衰减(使模型收敛更充分)

存在bug及其他问题私信：[email protected]

其他

该仓库可能存在bug，希望大家在使用过程中能及时反馈，或者留下一些代码修改意见。我们一起让它变得更好。

如果觉得有用清给我点star