Инференс изображений на предобученой модели

Перед началом работы

Зарегистрируйтесь в личном кабинете Cloud.ru.

Если вы уже зарегистрированы, войдите под своей учетной записью.

Шаг 1. Подготовьте среду

1. Убедитесь, что при создании ноутбука выбрали образ, который поддерживает CUDA.

2. Установите PyTorch и torchvision:

   pip install torch
   pip install torchvision
   
   

   Подробнее об установке PyTorch на официальном сайте.

3. Проверьте доступность GPU:

   import torch
   
   # Проверим доступность GPU
   cuda_available = torch.cuda.is_available()
   print(f"CUDA доступен: {cuda_available}")
   
   # Если GPU доступен — выведем информацию о доступности, количестве и названии GPU
   if cuda_available:
       print(f"Количество доступных GPU: {torch.cuda.device_count()}")
       print(f"Название GPU: {torch.cuda.get_device_name(0)}")
       device = torch.device("cuda")
   else:
       print("Используется CPU")
       device = torch.device("cpu")
   
   

4. Для начала работы с моделью и изображениями необходимо подключить следующие библиотеки:

   • torchvision — позволяет использовать предобученную модель ResNet18 и инструменты для преобразования изображений.

   • PIL — используется для работы с изображениями в формате PIL.

   • Requests — позволяет загружать изображения из интернета.

   • BytesIO из модуля io — используется для создания файлового объекта в памяти, который можно использовать для работы с данными, как если бы это был настоящий файл.

   • json и urllib.request — могут понадобиться для дополнительной работы с данными, например, для их сериализации или загрузки из интернета.

   Импортируйте библиотеки для работы:

   from torchvision import models, transforms
   from PIL import Image
   import requests
   from io import BytesIO
   import json
   import urllib.request
   
   

Шаг 2. Используйте предобученную модель для инференса

1. Загрузите изображение из интернета:

   url = "https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcSlSaRKDvkuC2_Qyfnt8jMDmZzphQbIrz-TSg&s"
   response = requests.get(url)
   img = Image.open(BytesIO(response.content)).convert('RGB')
   
   

   В результате мы получим объект изображения в формате RGB, готовый для дальнейшей обработки и анализа моделью.

2. Загрузите модель ResNet18 и переведите ее в режим инференса:

   model = models.resnet18(pretrained=True)
   model = model.to(device)
   model.eval()
   
   

   В результате мы загрузим модель, переведем ее в режим инференса и переместим на выбранное устройство GPU, если оно доступно, или CPU. Это необходимо для того, чтобы модель была готова к обработке изображений и выдавала предсказания.

Шаг 3. Подготовьте изображение и получите топ-5 предсказаний

1. Выполните преобразование изображения для инференса:

   preprocess = transforms.Compose([
       transforms.Resize(256),
       transforms.CenterCrop(224),
       transforms.ToTensor()
   ])
   
   

   Перед тем как подавать изображение на вход модели, необходимо преобразовать его в нужный формат. В данном случае мы применили три преобразования: изменение размера изображения до 256 пикселей, центрированный обрез до размера 224x224 пикселей и преобразование в tensor. Это необходимы для того, чтобы изображение соответствовало требованиям модели и могло быть обработано корректно.

2. Загрузите имена классов:

   url = "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
   urllib.request.urlretrieve(url, "imagenet_classes.txt")
   
   with open("imagenet_classes.txt") as f:
       class_names = [line.strip() for line in f.readlines()]
   
   

   В результате мы загрузили файл с именами классов ImageNet, которые используются в предобученной модели ResNet18. Это нужно для интерпретации результатов работы модели и понимания какие классы она может предсказывать.

3. Подготовьте изображение и выведите топ-5 предсказаний:

   # Подготовим изображение
   input_tensor = preprocess(img).unsqueeze(0).to(device)
   
   # Выполним инференс
   with torch.no_grad():
      output = model(input_tensor)
   
   # Прогоним через Softmax для получения вероятностей
   probabilities = torch.nn.functional.softmax(output[0], dim=0)
   
   # Получим топ-5 предсказаний
   top5_prob, top5_idx = torch.topk(probabilities, 5)
   
   # Выведем топ-5 предсказаний
   print("Топ-5 предсказаний:")
   for i in range(top5_prob.size(0)):
      class_name = class_names[top5_idx[i]]
      prob = top5_prob[i].item()
      print(f"{i+1}: {class_name} ({prob:.4f})")
   
   

   В результате мы получили Топ-5 предсказаний:

   Топ-5 предсказаний:
   1: tabby (0.5923)
   2: tiger cat (0.2160)
   3: Egyptian cat (0.1611)
   4: paper towel (0.0057)
   5: plastic bag (0.0037)