ResNet18开源模型体验:1块钱玩转图像分类,小白友好
1. 为什么选择ResNet18?
图像分类是AI领域最基础也最实用的技术之一。想象一下,如果你能教会电脑自动识别照片里的猫狗、花草、甚至不同品牌的汽车,是不是很酷?ResNet18就是这样一个帮你实现这个梦想的"小助手"。
作为中学生科技节的展示项目,ResNet18有三大优势:
- 轻量高效:相比动辄需要高端显卡的大模型,它能在普通电脑上运行
- 学习友好:模型结构清晰,特别适合理解深度学习基本原理
- 成本低廉:使用CSDN算力平台,最低1元就能体验完整流程
2. 5分钟快速部署
2.1 环境准备
你只需要: 1. 注册CSDN账号(已有可跳过) 2. 准备10-20张测试图片(手机随手拍即可) 3. 零花钱充值1元到账户
2.2 一键启动镜像
在CSDN算力平台搜索"ResNet18"镜像,选择PyTorch基础环境版本。点击"立即创建",系统会自动配置好所有依赖环境。
# 镜像已预装以下组件: - Python 3.8 - PyTorch 1.12 - torchvision 0.13 - 预训练好的ResNet18模型2.3 上传测试图片
将手机拍摄的图片通过网页端上传到/data目录。建议先准备两类图片(比如猫和狗),每类5-10张,这样效果最直观。
3. 实战图像分类
3.1 运行示例代码
新建Python文件demo.py,粘贴以下代码:
import torch from torchvision import transforms from PIL import Image # 加载预训练模型 model = torch.hub.load('pytorch/vision', 'resnet18', pretrained=True) model.eval() # 图像预处理 preprocess = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) # 加载测试图片 input_image = Image.open("/data/your_image.jpg") input_tensor = preprocess(input_image) input_batch = input_tensor.unsqueeze(0) # 使用GPU加速(如果有) if torch.cuda.is_available(): input_batch = input_batch.to('cuda') model.to('cuda') # 预测分类 with torch.no_grad(): output = model(input_batch) # 输出结果 probabilities = torch.nn.functional.softmax(output[0], dim=0) print(f"预测结果:{probabilities.argmax().item()}类,置信度:{probabilities.max().item():.2f}")3.2 查看分类结果
运行代码后,你会看到类似输出:
预测结果:282类,置信度:0.87这个数字对应ImageNet的类别索引。想知道具体是什么类别?可以添加这行代码:
# 在文件开头添加 import json with open("imagenet_class_index.json") as f: labels = json.load(f) # 替换最后的print语句为 pred_idx = probabilities.argmax().item() print(f"预测类别:{labels[str(pred_idx)][1]},置信度:{probabilities.max().item():.2f}")现在输出会变成:
预测类别:tiger cat,置信度:0.874. 自定义你的分类器
4.1 准备自己的数据集
- 新建
/data/train和/data/test文件夹 - 在每个文件夹内按类别建子文件夹(如
/data/train/cat,/data/train/dog) - 每类至少准备20张图片(手机拍摄即可)
4.2 微调模型
新建train.py文件,使用以下代码进行迁移学习:
import torch import torchvision from torch import nn, optim from torchvision import transforms, datasets # 数据预处理 transform = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 加载数据集 train_set = datasets.ImageFolder('/data/train', transform=transform) train_loader = torch.utils.data.DataLoader(train_set, batch_size=4, shuffle=True) # 修改模型最后一层 model = torchvision.models.resnet18(pretrained=True) num_features = model.fc.in_features model.fc = nn.Linear(num_features, 2) # 2分类任务 # 训练配置 criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9) # 训练循环 for epoch in range(5): # 训练5轮 running_loss = 0.0 for i, data in enumerate(train_loader, 0): inputs, labels = data optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() print(f'Epoch {epoch+1} loss: {running_loss/len(train_loader):.3f}') print('训练完成!') torch.save(model.state_dict(), 'my_model.pth')5. 常见问题解决
5.1 内存不足怎么办?
如果遇到CUDA内存错误,可以: - 减小batch_size(改为2或1) - 使用torch.cuda.empty_cache()清理缓存 - 选择更小的图片尺寸(如改为128x128)
5.2 预测结果不准?
尝试以下改进: 1. 增加每类图片数量(至少20张) 2. 调整学习率(lr参数,尝试0.01或0.0001) 3. 增加训练轮次(epochs参数)
5.3 如何保存/加载模型?
保存训练好的模型:
torch.save(model.state_dict(), 'my_model.pth')加载模型进行预测:
model.load_state_dict(torch.load('my_model.pth')) model.eval()6. 总结
通过这次实践,我们完成了:
- 零成本体验:用1元预算在CSDN平台完成AI模型部署
- 完整流程:从模型加载到自定义训练的全过程实践
- 实用技巧:掌握了图像分类的关键参数调整方法
- 创意扩展:学会了如何让模型认识你自己的图片类别
现在就可以上传你的照片,看看ResNet18能不能认出你的宠物或者收藏的玩具!实测下来,即使是2015年的模型,在今天的常见物体识别上依然表现不错。
💡获取更多AI镜像
想探索更多AI镜像和应用场景?访问 CSDN星图镜像广场,提供丰富的预置镜像,覆盖大模型推理、图像生成、视频生成、模型微调等多个领域,支持一键部署。