Создание полного, но быстрого конвейера компьютерного зрения для создания прототипа продукта или в качестве строительного блока для создания более сложной системы становится все более важным. Здесь мы собираемся обсудить, как вы можете сделать это за 5 минут или меньше, используя популярную библиотеку TorchVision.

Таблица содержания

  1. Готовиться
  2. Подготовка набора данных
  3. Конструкция модели
  4. Оценка модели
  5. Заключение

Готовиться

Первым делом вы должны установить библиотеки PyTorch и TorchVision. Рекомендуется делать это в виртуальной среде conda.

  1. Перейдите на страницу установки PyTorch, скопируйте и вставьте написанную командную строку в свой терминал, например:
conda install pytorch torchvision torchaudio -c pytorch

2. Убедитесь, что установка завершена, выполнив следующую команду в вашем терминале. Если ничего не выводится, значит, вы завершили этот раздел!

python -c "import torch"
python -c "import torchvision"

Подготовка набора данных

В этом руководстве мы собираемся использовать один из наиболее популярных наборов данных MNIST (набор изображений цифр от 0 до 9).

Во-первых, мы собираемся применить преобразование к набору данных, чтобы они хорошо вписались в нашу будущую модель AlexNet (например, изменение размера входного изображения до 224 и преобразование MNIST изображение в тензорный формат).

import torch
from torchvision import datasets, transforms, models
from torch.autograd import Variable
transform = transforms.Compose(
    [
        transforms.Resize(224),
        transforms.ToTensor(),
    ])

Затем мы загружаем как обучающий, так и тестовый набор, передавая вышеуказанный объект преобразования.

mnist_trainset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
mnist_testset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

Где корень указывает каталог, в котором данные будут загружены и сохранены.

Проверьте, завершена ли загрузка, узнав размер наборов данных.

len(mnist_trainset)
>>> 60000 
len(mnist_testset)
>>> 10000

Затем нам нужно создать загрузчик данных для пакетной обработки.

train_loader = torch.utils.data.DataLoader(mnist_trainset,
                                          batch_size=4,
                                          shuffle=True)
test_loader = torch.utils.data.DataLoader(mnist_testset,
                                          batch_size=4,
                                          shuffle=True)

Конструкция модели

Инициализировать модель также легко. Для нашего обсуждения мы будем использовать AlexNet с предварительно обученными весами (вы можете выбрать инициализацию модели со случайными весами и повторно обучить ее, используя эталонный набор данных, как указано выше, или используя свой собственный заранее определенный набор данных). Вы можете найти коллекцию моделей, поддерживаемых TorchVision, здесь.

alexnet = models.alexnet(pretrained= True)

Где вы затем можете найти детали моделей, запустив

alexnet

Самая первая строка выше указывает, что слой Conv2d в блоке (features) имеет размер входного канала 3 (RGB), тогда как наш набор данных MNIST имеет только один цветовой канал. :

alexnet.features[0] = torch.nn.Conv2d(1,64, kernel_size=(11,11), stride=(4,4), padding=(2,2))

Дальнейшие модификации тоже могут быть легко выполнены таким же образом!

Оценка модели

Теперь мы готовы обучать нашу модель.

  1. Определите функцию потерь и оптимизатор

В этом случае мы используем CrossEntropyLoss и Stochastic Gradient Descent (SGD) соответственно.

criterion=torch.nn.CrossEntropyLoss()
optimizer=torch.optim.SGD(alexnet.parameters(),lr=0.01,momentum=0.9)

2. Обучите и подтвердите

print ("Training begins...")
num_of_epochs = 3
for epoch in range(num_of_epochs):
    running_loss=0
    for i, data in enumerate(train_loader, 0):
        image, label = data
 
        image = Variable(image)
        label = Variable(label)
        optimizer.zero_grad()
 
        outputs = alexnet(image)
        loss = criterion(outputs, label)
 
        loss.backward()
        optimizer.step()
 
        running_loss += loss.data
 
        if i%100==99:
            print (f'[epoch {epoch + 1}, imgs {i + 1}] loss:  {running_loss/100}')
            running_loss = 0
print("Training finishes...")

Вышеупомянутые подпрограммы обучения и проверки являются стандартными (т. Е. Включают функции zero_grad (), backward () и step ()). Конечно, вы можете изменить гиперпараметры, включая эпоху обучения, скорость обучения и т. Д., Для оптимизации.

3. Вывод

Подобно обучению, мы можем выполнить вывод следующим образом:

alexnet.eval()
correct=0
total=0
for data in test_loader:
    images, labels = data
    outputs = alexnet(Variable(images))
    _ , predicted = torch.max(outputs, 1)
    total += labels.size(0)

Заключение

Вы только что построили непрерывный конвейер CV от сбора данных, преобразования, построения модели, обучения, проверки и вывода до логического вывода - и все это менее чем за пять минут! Я также продемонстрировал, как легко изменить вашу модель в соответствии с вашим собственным набором данных. Аналогичные корректировки могут быть внесены во всю модель.

Подпишитесь на мою информационную рассылку по электронной почте: https://tinyurl.com/2npw2fnz , где я регулярно резюмирую исследовательские работы по ИИ на простом английском языке и в красивой визуализации.