Ubuntu20.04下ViT图像分类模型从零部署教程1. 准备工作与环境配置在开始部署ViT图像分类模型之前我们需要确保系统环境准备就绪。Ubuntu20.04是一个稳定的起点但还需要一些额外的配置。首先更新系统包列表确保我们安装的是最新版本的软件包sudo apt update sudo apt upgrade -y安装Python环境管理工具推荐使用miniconda来创建独立的环境wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh bash Miniconda3-latest-Linux-x86_64.sh创建专门的Python环境来部署我们的模型conda create -n vit-env python3.8 conda activate vit-env2. 安装必要的依赖包ViT模型依赖于一些关键的Python库我们需要逐一安装它们。这些依赖包括深度学习框架和图像处理工具。安装PyTorch和相关的计算机视觉库pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu如果你有NVIDIA显卡并希望使用GPU加速可以安装CUDA版本的PyTorchpip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118安装其他必要的依赖包pip install transformers pillow numpy matplotlib requests验证安装是否成功import torch print(fPyTorch版本: {torch.__version__}) print(fCUDA是否可用: {torch.cuda.is_available()})3. 下载和配置ViT模型现在我们来获取ViT模型。我们将使用Hugging Face提供的预训练模型这是一个很好的起点。创建模型目录并下载所需文件from transformers import ViTForImageClassification, ViTFeatureExtractor import torch # 下载模型和特征提取器 model_name google/vit-base-patch16-224 model ViTForImageClassification.from_pretrained(model_name) feature_extractor ViTFeatureExtractor.from_pretrained(model_name) # 保存模型到本地 model.save_pretrained(./vit-model) feature_extractor.save_pretrained(./vit-model)为了确保模型能够正确识别日常物品我们还需要配置标签文件。ViT模型使用ImageNet的1000个类别标签import requests # 下载ImageNet标签 labels_url https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt response requests.get(labels_url) labels response.text.split(\n) # 保存标签文件 with open(imagenet_labels.txt, w) as f: for label in labels: f.write(f{label}\n)4. 编写图像分类推理代码有了模型和标签我们现在可以编写实际的分类代码了。这段代码将接收图像输入并返回分类结果。创建主要的推理脚本import torch from PIL import Image from transformers import ViTForImageClassification, ViTFeatureExtractor import requests from io import BytesIO class ViTImageClassifier: def __init__(self, model_path./vit-model): self.device torch.device(cuda if torch.cuda.is_available() else cpu) self.model ViTForImageClassification.from_pretrained(model_path) self.model.to(self.device) self.model.eval() self.feature_extractor ViTFeatureExtractor.from_pretrained(model_path) # 加载标签 with open(imagenet_labels.txt, r) as f: self.labels [line.strip() for line in f.readlines() if line.strip()] def classify_image(self, image_path): # 加载图像 if image_path.startswith(http): response requests.get(image_path) image Image.open(BytesIO(response.content)) else: image Image.open(image_path) # 预处理图像 inputs self.feature_extractor(imagesimage, return_tensorspt) inputs {k: v.to(self.device) for k, v in inputs.items()} # 进行推理 with torch.no_grad(): outputs self.model(**inputs) # 处理结果 probabilities torch.nn.functional.softmax(outputs.logits, dim-1) top5_prob, top5_indices torch.topk(probabilities, 5) results [] for i in range(5): label self.labels[top5_indices[0][i].item()] probability top5_prob[0][i].item() results.append((label, probability)) return results # 使用示例 if __name__ __main__: classifier ViTImageClassifier() # 测试本地图像 results classifier.classify_image(test_image.jpg) print(分类结果:) for label, prob in results: print(f{label}: {prob:.4f})5. 测试模型效果现在让我们测试一下部署的模型是否正常工作。我们将使用一些示例图像来验证分类效果。创建测试脚本def test_with_sample_images(): classifier ViTImageClassifier() # 测试图像URL可以使用任何在线图像 test_images [ https://images.unsplash.com/photo-1541963463532-d68292c34b19, # 书籍 https://images.unsplash.com/photo-1526170375885-4d8ecf77b99f, # 相机 https://images.unsplash.com/photo-1575936123452-b67c3203c357 # 猫 ] for i, img_url in enumerate(test_images): print(f\n测试图像 {i1}: {img_url}) try: results classifier.classify_image(img_url) print(前5个预测结果:) for j, (label, prob) in enumerate(results): print(f{j1}. {label}: {prob:.4f}) except Exception as e: print(f处理图像时出错: {e}) # 运行测试 test_with_sample_images()你也可以测试本地图像# 测试本地图像文件 local_results classifier.classify_image(/path/to/your/image.jpg) print(本地图像分类结果:) for label, prob in local_results: print(f{label}: {prob:.4f})6. 常见问题与解决方案在部署过程中可能会遇到一些问题这里列出了一些常见问题及其解决方法。问题1内存不足错误如果遇到内存不足的问题可以尝试减小批处理大小或使用更小的模型# 使用较小的模型变体 model_name google/vit-base-patch16-224 # 可以尝试其他变体问题2CUDA内存不足如果使用GPU时遇到内存问题# 减少批处理大小 inputs self.feature_extractor(imagesimage, return_tensorspt)问题3图像预处理问题确保图像格式正确# 检查图像模式并转换为RGB if image.mode ! RGB: image image.convert(RGB)问题4依赖版本冲突创建新的conda环境可以避免大多数依赖冲突# 如果遇到依赖问题创建全新环境 conda create -n vit-new python3.8 conda activate vit-new问题5模型下载缓慢可以设置镜像源加速下载import os os.environ[HF_ENDPOINT] https://hf-mirror.com7. 总结完成以上步骤后你应该已经在Ubuntu20.04系统上成功部署了ViT图像分类模型。整个过程从环境准备开始逐步安装了必要的依赖下载并配置了预训练模型最后编写了推理代码来测试模型效果。实际使用中这个模型能够识别1000种不同的物体类别覆盖了日常生活中大多数常见物品。你可以进一步扩展这个基础部署比如添加Web界面、实现批量处理功能或者针对特定领域进行模型微调。记得定期更新依赖包以获取性能改进和安全修复。如果遇到性能问题可以考虑使用模型量化或剪枝技术来优化推理速度。对于生产环境部署还需要添加适当的错误处理和日志记录机制。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。