使用PyTorch进行图像迁移学习

本示例解释了使用Label Studio和PyTorch进行计算机视觉的基础知识。 该模型采用了流行的ResNet图像分类器的迁移学习技术，并可针对您的数据进行微调。

你可以使用以下示例标注配置：

<View>
  <Image name="image_object" value="$image_url"/>
  <Choices name="image_classes" toName="image_object">
    <Choice value="Cat"/>
    <Choice value="Dog"/>
  </Choices>
</View>

创建模型脚本

如果使用Label Studio的ML SDK创建机器学习后端，请确保您的ML后端脚本执行以下操作：

• 继承创建的模型类自 label_studio_ml.LabelStudioMLBase
• Override the 2 methods:
  • predict()，接收输入任务并以Label Studio JSON格式输出预测结果。
  • fit()，接收annotations可迭代对象并返回包含已创建链接和资源的字典。该字典随后会通过self.train_output字段用于加载模型。

创建一个名为model.py的文件，其中包含准备好用于训练和推理的PyTorch模型。

首先，创建一个Dataset类，该类接收图像URL列表作为输入，并生成一批经过预处理的图像及其对应标签：

import torch
import torch.nn as nn
import torch.optim as optim
import time

from torch.utils.data import Dataset, DataLoader
from torchvision import datasets, models, transforms

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')


class ImageClassifierDataset(Dataset):
        
    def __init__(self, image_urls, image_classes):
        self.images = []
        self.labels = []
        
        self.classes = list(set(image_classes))
        self.class_to_label = {c: i for i, c in enumerate(self.classes)}
        
        self.image_size = 224
        self.transforms = transforms.Compose([
            transforms.Resize(self.image_size),
            transforms.CenterCrop(self.image_size),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        
        
        for image_url, image_class in zip(image_urls, image_classes):
            image = self._get_image_from_url(image_url)
            transformed_image = self.transforms(image)
            self.images.append(transformed_image)
            
            label = self.class_to_label[image_class]
            self.labels.append(label)
            
    def _get_image_from_url(self, url):
        pass
    
    def __getitem__(self, index):
        return self.images[index], self.labels[index]
    
    def __len__(self):
        return len(self.images)

接下来，为预训练的ResNet模型创建一个简单的封装器：

class ImageClassifier(object):
    
    def __init__(self, num_classes):
        self.model = models.resnet18(pretrained=True)
        num_ftrs = self.model.fc.in_features
        self.model.fc = nn.Linear(num_ftrs, num_classes)
        
        self.model = self.model.to(device)
        
        self.criterion = nn.CrossEntropyLoss()
        self.optimizer = optim.SGD(self.model.parameters(), lr=0.001, momentum=0.9)
        
        # Decay LR by a factor of 0.1 every 7 epochs
        self.scheduler = optim.lr_scheduler.StepLR(self.optimizer, step_size=7, gamma=0.1)
    
    def save(self, path):
        torch.save(self.model.state_dict(), path)
    
    def load(self, path):
        self.model.load_state_dict(torch.load(path))
        self.model.eval()
        
    def train(self, dataloader, num_epochs=25):
        since = time.time()

        self.model.train()
        for epoch in range(num_epochs):
            print('Epoch {}/{}'.format(epoch, num_epochs - 1))
            print('-' * 10)
            
            running_loss = 0.0
            running_corrects = 0    
            # Iterate over data.
            for inputs, labels in dataloader:
                inputs = inputs.to(device)
                labels = labels.to(device)

                self.optimizer.zero_grad()
                outputs = self.model(inputs)
                _, preds = torch.max(outputs, 1)
                loss = self.criterion(outputs, labels)
                loss.backward()
                self.optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
                self.scheduler.step()

            epoch_loss = running_loss / len(dataloader.dataset)
            epoch_acc = running_corrects.double() / len(dataloader.dataset)

            print('Train Loss: {:.4f} Acc: {:.4f}'.format(epoch_loss, epoch_acc))

        print()
    
        time_elapsed = time.time() - since
        print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))

        return self.model

最后，重写API方法：

from label_studio_ml.model import LabelStudioMLBase

class ImageClassifierAPI(LabelStudioMLBase):
    
    def __init__(self, **kwargs):
        self.model = ImageClassifier(resources['num_classes'])
        self.model.load(resources['model_path'])
        self.labels = resources['labels']
        
    def predict(self, tasks, **kwargs):
        pass

    def fit(self, completions, **kwargs):
        pass

创建ML后端配置与脚本

Label Studio可以自动创建运行ML后端所需的所有配置和脚本，这些配置和脚本来自您新创建的模型。

调用您的机器学习后端 my_backend，并通过命令行初始化机器学习后端目录 ./my_backend：

label-studio-ml init my_backend

最后一条命令会获取您的脚本./model.py并在同级目录下创建./my_backend文件夹，复制启动ML后端所需的各种配置和脚本，适用于开发或生产模式。

启动ML后端服务器

开发模式

在开发模式下，训练和推理在同一个进程中完成，因此服务器在模型训练期间不会响应传入的预测请求。

要在Flask开发模式下启动ML后端服务器，请从命令行运行以下命令：

label-studio-ml start my_backend

服务器已启动在 http://localhost:9090 并在控制台输出日志。

生产模式

生产模式由Redis服务器和RQ作业提供支持，负责后台训练流程。这意味着您可以启动模型训练，同时继续从当前模型状态请求预测。当模型完成训练后，新版本模型会自动更新。

对于生产环境，请确保您的系统已安装Docker和docker-compose。然后在命令行中运行以下命令：

cd my_backend/
docker-compose up

您可以在my_backend/logs/uwsgi.log中查看运行时日志，在my_backend/logs/rq.log中查看RQ训练日志

在Label Studio中使用机器学习后端

初始化并启动一个新的Label Studio项目，连接到正在运行的ML后端：

label-studio start my_project --init --ml-backends http://localhost:9090

获取预测结果

您应该在标注界面中看到模型预测结果。参见使用Label Studio设置机器学习。

模型训练

通过点击项目设置中机器学习页面的Start training按钮手动触发模型训练，或使用API调用：

curl -X POST http://localhost:8080/api/models/train