FlagAI新增视觉模型！支持一键训练Swin Transformer，还有更多功能更新

import os
import torch
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from flagai.trainer import Trainer
from flagai.auto_model import AutoLoader

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
data_path = "./imagenet2012/"

# use DDP for training by 4 gpus.
trainer = Trainer(env_type="pytorchDDP",
                  epochs=10,
                  experiment_name="swinv2_imagenet_ddp",
                  batch_size=32,
                  weight_decay=1e-3,
                  warm_up=0.1,
                  lr=5e-5,
                  save_interval=100,
                  eval_interval=100,
                  log_interval=10,
                  num_gpus=4,
                  hostfile="./hostfile",
                  training_script="training_swinv2.py"
                  )

# swinv2 model_name support:
# 1. swinv2-base-patch4-window16-256,
# 2. swinv2-small-patch4-window16-256,
# 3. swinv2-base-patch4-window8-256
loader = AutoLoader(task_name="classification",
                    model_name="swinv2-base-patch4-window8-256",
                    num_classes=1000)
model = loader.get_model()

# build imagenet dataset
def build_dataset(root):
    traindir = os.path.join(root, 'train')
    valdir = os.path.join(root, 'val')
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

    train_dataset = datasets.ImageFolder(
        traindir,
        transforms.Compose([
            transforms.RandomResizedCrop(256),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize
        ])
    )
    val_dataset = datasets.ImageFolder(
        valdir,
        transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            normalize
        ])
    )
    return train_dataset, val_dataset

def collate_fn(batch):
    images = torch.stack([b[0] for b in batch])
    if trainer.fp16:
        images = images.half()
    labels = [b[1] for b in batch]
    labels = torch.tensor(labels).long()
    return {"images": images, "labels": labels}

def top1_acc(pred, labels, **kwargs):
    pred = pred.argmax(dim=1)
    top1_acc = pred.eq(labels).sum().item() / len(pred)
    return top1_acc

if __name__ == '__main__':

    print("building imagenet dataset......")
    train_dataset, val_dataset = build_dataset(root=data_path)
    print("training......")
    optimizer = torch.optim.Adam(model.parameters(), lr=5e-5)
    trainer.train(model,
                  train_dataset=train_dataset,
                  valid_dataset=val_dataset,
                  collate_fn=collate_fn,
                  optimizer=optimizer,
                  metric_methods=[["top1_acc", top1_acc]],
                  find_unused_parameters=False)

import torch
import os
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
import torchvision.datasets as datasets
from tqdm import tqdm
from flagai.auto_model.auto_loader import AutoLoader

data_path = "./imagenet2012/"

# swinv2 model_name support:
# 1. swinv2-base-patch4-window16-256,
# 2. swinv2-small-patch4-window16-256,
# 3. swinv2-base-patch4-window8-256
loader = AutoLoader(task_name="classification",
                    model_name="swinv2-small-patch4-window16-256")
model = loader.get_model()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model.eval()
model.to(device)

# imagenet loader
def data_loader(root, batch_size=256, workers=1):
    valdir = os.path.join(root, 'val')
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

    val_dataset = datasets.ImageFolder(
        valdir,
        transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            normalize
        ])
    )

    val_loader = DataLoader(val_dataset,
                            batch_size=batch_size,
                            shuffle=False,
                            num_workers=workers
                            )

    return val_loader

@torch.no_grad()
def test(model,data_loader):
    model.eval()
    top1_acc = 0.0
    top5_acc = 0.0

    for step, (inputs, labels) in tqdm(enumerate(data_loader), total=len(data_loader)):
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)["logits"]
        _, top1_preds = outputs.max(1)
        top1_acc += top1_preds.eq(labels).sum().item()

        top5_pred = outputs.topk(5, 1, True)[1]
        top5_acc += top5_pred.eq(labels.view(-1, 1).expand_as(top5_pred).to(device)).sum().item()

    print(
        "test_top1_acc [{top1_acc}], test_top5_acc [{top5_acc}] \n".format(
            top1_acc=top1_acc/len(data_loader.dataset),
            top5_acc=top5_acc/len(data_loader.dataset),
        )
    )
if __name__ == '__main__':

    val_loader = data_loader(data_path, batch_size=8, workers=8)
    test(model, val_loader)