Pytorch Accelerate multi-GPU Training Inference

Pytorch Accelerate multi-GPU Training Inference

• enhance the training speed in around 2 times
• easy to use the mix way in training to perform the acceleration

#GPU	fp16	Batch size per GPU	Seconds per epoch
1	no	256	60
2	no	256	33
2	no	128	41
2	yes	128	35

Use Accelerate change the code of single GPU to implement the training and inference of multi-GPUs

Accelerate the parameters and mix-training

multi-GPU to print the information of main process

multi-GPU training and set the right batch size

Single-GPU

import pandas as pd
import numpy as np
import cv2

import torch
import torch.nn as nn
from torch.optim import Adam
from torch.utils.data import Dataset, DataLoader, Subset
from torch.nn import CrossEntropyLoss

import albumentations as A
from albumentations.pytorch import ToTensorV2

import timm
from tqdm import tqdm


def get_transform(image_size, train=True):
    if train:
        return A.Compose([
                    A.HorizontalFlip(p=0.5),
                    A.VerticalFlip(p=0.5),
                    A.RandomBrightnessContrast(p=0.2),
                    A.Resize(image_size, image_size, interpolation=cv2.INTER_LANCZOS4),
                    A.Normalize(0.1310, 0.30854),
                    ToTensorV2(),
                ])
    else:
        return A.Compose([
                    A.Resize(image_size, image_size, interpolation=cv2.INTER_LANCZOS4),
                    A.Normalize(0.1310, 0.30854),
                    ToTensorV2(),
                ])
    
    
class MiniDataSet(Dataset):
    
    def __init__(self, images, labels=None, transform=None):
        self.images = images.astype("float32")
        self.labels = labels
        self.transform = transform
    
    def __len__(self):
        return len(self.images)
    
    def __getitem__(self, idx):
        ret = {}
        img = self.images[idx]
        
        if self.transform is not None:
            img = self.transform(image=img)["image"]
        ret["image"] = img
        
        if self.labels is not None:
            ret["label"] = self.labels[idx]
        
        return ret
    
    
class TimmModel(nn.Module):
    
    def __init__(self, backbone, num_class, pretrained=False):
        
        super().__init__()
        self.model = timm.create_model(backbone, pretrained=pretrained, in_chans=1, num_classes=num_class)
    
    def forward(self, image):
        logit = self.model(image)
        return logit


def train_fn(args, model, optimizer, dataloader):
    model.to(args.device)
    model.train()
    train_loss = []
    loss_fn = torch.nn.CrossEntropyLoss()
    
    for batch in tqdm(dataloader):
        logits = model(batch["image"].to(args.device))
        optimizer.zero_grad()
        loss = loss_fn(logits, batch["label"].to(args.device))
        loss.backward()
        optimizer.step()
        train_loss.append(loss.item())
    
    return np.mean(train_loss)


@torch.no_grad()
def predict_fn(args, model, dataloader):
    model.to(args.device)
    model.eval()
    predictions = []
    
    for step, batch in enumerate(dataloader):
        output = model(batch["image"].to(args.device))
        prediction = torch.argmax(output, 1)
        predictions.append(prediction.cpu().numpy())
    
    predictions = np.concatenate(predictions, axis=0)
    return predictions


def fit_model(args, model, optimizer, train_dl, val_dl):
    best_score = 0.
    for ep in range(args.ep):
        train_loss = train_fn(args, model, optimizer, train_dl)
        val_pred = predict_fn(args, model, val_dl)
        val_acc = np.mean(val_pred == val_dl.dataset.labels)
        print(f"Epoch {ep+1}, train loss {train_loss:.4f}, val acc {val_acc:.4f}")
        if val_acc > best_score:
            best_score = val_acc
            torch.save(model.state_dict(), "model.bin")
    model.load_state_dict(torch.load("model.bin"))
    return model
    
    

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--ep", default=5, type=int)
    parser.add_argument("--lr", default=0.00005, type=float)
    parser.add_argument("--bs", default=256, type=int)
    parser.add_argument("--device", default=0, type=int)
    parser.add_argument("--model", default="convnext_small")
    parser.add_argument("--image_size", default=56, type=int)
    args = parser.parse_args()
    
    train = pd.read_csv("digit-recognizer/train.csv")
    train_images = train.iloc[:, 1:].values.reshape(-1, 28, 28)
    train_labels = train.iloc[:, 0].values
    
    test = pd.read_csv("digit-recognizer/test.csv")
    test_images = test.values.reshape(-1, 28, 28)
    
    submission = pd.read_csv("digit-recognizer/sample_submission.csv")

    
    train_transform = get_transform(args.image_size, True)
    valid_transform = get_transform(args.image_size, False)
    train_ds = MiniDataSet(train_images[:40000], train_labels[:40000], train_transform)
    val_ds = MiniDataSet(train_images[40000:], train_labels[40000:], valid_transform)
    test_ds = MiniDataSet(test_images, transform=valid_transform)
    
    train_dl = DataLoader(
    train_ds, 
    batch_size=args.bs, 
    num_workers=2, 
    shuffle=True, 
    drop_last=True)

    val_dl = DataLoader(
        val_ds, 
        batch_size=args.bs * 2, 
        num_workers=2, 
        shuffle=False, 
        drop_last=False)

    test_dl = DataLoader(
        test_ds, 
        batch_size=args.bs * 2, 
        num_workers=2, 
        shuffle=False, 
        drop_last=False)


    train_ds = MiniDataSet(train_images[:40000], train_labels[:40000], train_transform)
    val_ds = MiniDataSet(train_images[40000:], train_labels[40000:], valid_transform)
    test_ds = MiniDataSet(test_images, transform=valid_transform)
    
    model = TimmModel(args.model, 10, pretrained=True)
    optimizer = Adam(model.parameters(), lr=args.lr)
    model = fit_model(args, model, optimizer, train_dl, val_dl)
    
    test_pred = predict_fn(args, model, test_dl)
    submission["Label"] = test_pred
    submission.to_csv("./submission.csv", index=False)

Modify code for Multi-GPUs

warnings.filterwarnings("ignore")

def get_transform(image_size, train=True):
    if train:
        return A.Compose([
                    A.HorizontalFlip(p=0.5),
                    A.VerticalFlip(p=0.5),
                    A.RandomBrightnessContrast(p=0.2),
                    A.Resize(image_size, image_size, interpolation=cv2.INTER_LANCZOS4),
                    A.Normalize(0.1310, 0.30854),
                    ToTensorV2(),
                ])
    else:
        return A.Compose([
                    A.Resize(image_size, image_size, interpolation=cv2.INTER_LANCZOS4),
                    A.Normalize(0.1310, 0.30854),
                    ToTensorV2(),
                ])
    
    
class MiniDataSet(Dataset):
    
    def __init__(self, images, labels=None, transform=None):
        self.images = images.astype("float32")
        self.labels = labels
        self.transform = transform
    
    def __len__(self):
        return len(self.images)
    
    def __getitem__(self, idx):
        ret = {}
        img = self.images[idx]
        
        if self.transform is not None:
            img = self.transform(image=img)["image"]
        ret["image"] = img
        
        if self.labels is not None:
            ret["label"] = self.labels[idx]
        
        return ret
    
    
class TimmModel(nn.Module):
    
    def __init__(self, backbone, num_class, pretrained=False):
        
        super().__init__()
        self.model = timm.create_model(backbone, pretrained=pretrained, in_chans=1, num_classes=num_class)
    
    def forward(self, image):
        logit = self.model(image)
        return logit
    
    
def train_fn(args, model, optimizer, dataloader):
#     model.to(args.device)
    model.train()
    train_loss = []
    loss_fn = torch.nn.CrossEntropyLoss()

#     增加disable=not args.accelerator.is_main_process，只在主进程显示进度条，避免重复显示
#     for batch in tqdm(dataloader):
    for batch in tqdm(dataloader, disable=not args.accelerator.is_main_process):
#         logits = model(batch["image"].to(args.device))
        logits = model(batch["image"])
        optimizer.zero_grad()
#         loss = loss_fn(logits, batch["label"].to(args.device))
        loss = loss_fn(logits, batch["label"])
#         loss.backward 修改为 accelerator.backward(loss)
#         loss.backward()
        args.accelerator.backward(loss)
        optimizer.step()
        train_loss.append(loss.item())
    
    return np.mean(train_loss)


@torch.no_grad()
def predict_fn(args, model, dataloader):
#     删除 to(device)
#     model.to(args.device)
    model.eval()
    predictions = []
    
    for step, batch in enumerate(dataloader):
#         output = model(batch["image"].to(args.device))
        output = model(batch["image"])
        prediction = torch.argmax(output, 1)
#         使用accelerator.gather_for_metrics(prediction)汇总多张GPU预测结果
        prediction = args.accelerator.gather_for_metrics(prediction)
        predictions.append(prediction.cpu().numpy())
    
    predictions = np.concatenate(predictions, axis=0)
    return predictions


# def train_fn(args, model, optimizer, dataloader):
#     # 删除 to(device)
#     # loss.backward 修改为 accelerator.backward(loss)
#     # 增加disable=not args.accelerator.is_main_process，只在主进程显示进度条，避免重复显示
#     model.train()
#     train_loss = []
#     loss_fn = torch.nn.CrossEntropyLoss()
    
#     for batch in tqdm(dataloader, disable=not args.accelerator.is_main_process):
#         logits = model(batch["image"])
#         optimizer.zero_grad()
#         loss = loss_fn(logits, batch["label"])
#         args.accelerator.backward(loss)
#         optimizer.step()
#         train_loss.append(loss.item())
    
#     return np.mean(train_loss)


# @torch.no_grad()
# def predict_fn(args, model, dataloader):
#     # 删除 to(device)
#     # 使用accelerator.gather_for_metrics(prediction)汇总多张GPU预测结果
#     model.eval()
#     predictions = []
    
#     for step, batch in enumerate(dataloader):
#         output = model(batch["image"])
#         prediction = torch.argmax(output, 1)
#         prediction = args.accelerator.gather_for_metrics(prediction)
#         predictions.append(prediction.cpu().numpy())
    
#     predictions = np.concatenate(predictions, axis=0)
#     return predictions


def fit_model(args, model, optimizer, train_dl, val_dl):
    best_score = 0.
    for ep in range(args.ep):
        train_loss = train_fn(args, model, optimizer, train_dl)
        val_pred = predict_fn(args, model, val_dl)
        val_acc = np.mean(val_pred == val_dl.dataset.labels)
        if args.accelerator.is_main_process:
            print(f"Epoch {ep+1}, train loss {train_loss:.4f}, val acc {val_acc:.4f}")
        if val_acc > best_score:
            best_score = val_acc
            torch.save(model.state_dict(), "model.bin")
    model.load_state_dict(torch.load("model.bin"))
    return model
    
    

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--ep", default=5, type=int)
    parser.add_argument("--lr", default=0.00005, type=float)
    parser.add_argument("--bs", default=256, type=int)
    parser.add_argument("--device", default=0, type=int)
    parser.add_argument("--model", default="convnext_small")
    parser.add_argument("--image_size", default=56, type=int)
    args = parser.parse_args()
    
    
    train = pd.read_csv("digit-recognizer/train.csv")
    train_images = train.iloc[:, 1:].values.reshape(-1, 28, 28)
    train_labels = train.iloc[:, 0].values
    
    test = pd.read_csv("digit-recognizer/test.csv")
    test_images = test.values.reshape(-1, 28, 28)
    
    submission = pd.read_csv("digit-recognizer/sample_submission.csv")

    
    train_transform = get_transform(args.image_size, True)
    valid_transform = get_transform(args.image_size, False)
    train_ds = MiniDataSet(train_images[:40000], train_labels[:40000], train_transform)
    val_ds = MiniDataSet(train_images[40000:], train_labels[40000:], valid_transform)
    test_ds = MiniDataSet(test_images, transform=valid_transform)
    
    train_dl = DataLoader(
    train_ds, 
    batch_size=args.bs, 
    num_workers=2, 
    shuffle=True, 
    drop_last=True)

    val_dl = DataLoader(
        val_ds, 
        batch_size=args.bs * 2, 
        num_workers=2, 
        shuffle=False, 
        drop_last=False)

    test_dl = DataLoader(
        test_ds, 
        batch_size=args.bs * 2, 
        num_workers=2, 
        shuffle=False, 
        drop_last=False)
    
    model = TimmModel(args.model, 10, pretrained=True)
    optimizer = Adam(model.parameters(), lr=args.lr)
    
    # 初始化Accelerator
    accelerator = Accelerator()
    
    # 多GPU训练准备
    model, optimizer, train_dl, val_dl, test_dl = accelerator.prepare(model, optimizer, train_dl, val_dl, test_dl)
    args.accelerator = accelerator
    
    model = fit_model(args, model, optimizer, train_dl, val_dl)
    
    test_pred = predict_fn(args, model, test_dl)
    if accelerator.is_local_main_process:
        submission["Label"] = test_pred
        submission.to_csv("./submission.csv", index=False)