侧边栏壁纸
博主头像
小鱼吃猫博客博主等级

你所热爱的,便是你的生活。

  • 累计撰写 114 篇文章
  • 累计创建 47 个标签
  • 累计收到 9 条评论

目 录CONTENT

文章目录

图像分类实战——阿猫阿狗的识别

小鱼吃猫
2023-08-31 / 0 评论 / 4 点赞 / 60 阅读 / 3923 字

概述

做一个阿猫阿狗图片分类的功能,数据集和模型均来自Huggingface 数据集:Bingsu/Cat_and_Dog 模型:microsoft/resnet-50

数据加载

数据集共包含1w个数据,分为训练集8000,测试集2000。原始数据集在Hugginface上,需要联网下载,可以下载离线数据集然后从本地加载(模型和数据在文末的地址)。 数据示例 小鱼吃猫博客

image_data = load_dataset('Bingsu/Cat_and_Dog')

数据处理

image_processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50")

def image_detail(examples):
    images =[Image.open(io.BytesIO(image['bytes'])) for image in examples['image']]
    inputs = image_processor(images, return_tensors="pt")
    inputs['labels']=examples['labels']
    return inputs

image_data_pt=image_data.map(image_detail,batched=True,remove_columns=image_data["train"].column_names)

加载模型

id2label={0:'cat', 1:'dog'}
label2id={'cat':0, 'dog':1}
model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-50",num_labels=2,id2label=id2label,label2id=label2id,ignore_mismatched_sizes=True)

评估函数

如果你不想在训练过程中对模型进行评估,可以不写。只需要在最后做一次评估就行。

import evaluate

acc_metric = evaluate.load("accuracy")
f1_metirc = evaluate.load("f1",average='micro')
def eval_metric(eval_predict):
    predictions, labels = eval_predict
    predictions = predictions.argmax(axis=-1)
    acc = acc_metric.compute(predictions=predictions, references=labels)
    f1 = f1_metirc.compute(predictions=predictions, references=labels,average='micro')
    acc.update(f1)
    return acc

训练参数

train_args =TrainingArguments(output_dir="./checkpoints",      # 输出文件夹
				   per_device_train_batch_size=512,   # 训练时的batch_size
				   per_device_eval_batch_size=1,    # 验证时的batch_size
				   num_train_epochs=1,              # 训练轮数
				   logging_steps=10,                # log 打印的频率
				   evaluation_strategy="epoch",     # 评估策略
				   save_strategy="epoch",           # 保存策略
				   learning_rate=2e-5,              # 学习率
				   metric_for_best_model="f1",      # 设定评估指标
				   load_best_model_at_end=True,    # 训练完成后加载最优模型  
                    )    

训练器

trainer = Trainer(model=model, #模型
                  args=train_args,  # 训练参数
                  train_dataset=image_data_pt["train"],  # 训练集
                  eval_dataset=image_data_pt["test"], # 验证集合
                  compute_metrics=eval_metric) # 评估函数

模型训练

配置好所有参数后,使用训练器开启训练

trainer.train()

我在训练过程中的评估结果

EpochTraining LossValidation LossAccuracyF1
1No log0.6768350.6620000.662000

模型推理

from transformers import pipeline
pipe = pipeline("image-classification", model=model,image_processor=image_processor,device=0)

train_image = image_data['test']['image']
show_examples=train_image[:15]
for image in show_examples:
    _image = Image.open(io.BytesIO(image['bytes']))
    _image=_image.resize((100,100))
    res=pipe(_image)
    result = sorted(res,key=lambda x:x['score'],reverse=True)[0]['label']
    _image.show(title='result')
    print(result)
    print('------------')

我只取了测试集中的前15条进行推理预测,预测结果如下: dog_cat.png 可以看得出来,第一个小猫就被识别错了,哈哈,效果一般般啦。原因有很多,跟数据集的大小,训练的次数,学习率等多种因素有关,感兴趣的小伙伴可以自行尝试下。

4

评论区