低精度优化器 --> 高精度优化器状态 --> 结合梯度更新参数 --> 重新量化回低精度参数

pip install bitsandbytes-cuda111

import bitsandbytes as bnb
import torch
from torch import nn
import torchvision
import torchvision.transforms as transforms
import numpy as np
import random
from torchvision import models
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

torch.manual_seed(0)
random.seed(0)
np.random.seed(0)

batch_size=128
epochs=200

# we will apply the same transforms as described in the paper
train_transform=transforms.Compose(
[transforms.RandomHorizontalFlip(),
     transforms.ToTensor(),
     transforms.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261))])

val_transform=transforms.Compose(
[transforms.ToTensor(),
     transforms.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261))])


trainset=torchvision.datasets.CIFAR10(root='https://www.zhihu.com/question/cifar10', train=True,
                                        download=True, transform=train_transform)

train_dataloader=torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                          shuffle=True, num_workers=8)

testset=torchvision.datasets.CIFAR10(root='https://www.zhihu.com/question/cifar10', train=False,
                                       download=True, transform=val_transform)

val_dataloader=torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                         shuffle=False, num_workers=8)

classes=('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

iter=0

model=models.resnet34(num_classes=10).cuda()
criterion=nn.CrossEntropyLoss()

#Uncomment to use the int8 Optimizer
#logger=SummaryWriter(log_dir="log_adam_int8")
#optimizer=bnb.optim.Adam8bit(model.parameters(), lr=0.01, weight_decay=0.0005)

logger=SummaryWriter(log_dir="log_adam")
optimizer=optim.Adam(model.parameters(), lr=0.01, weight_decay=0.0005)

loss_plot_name='loss'
train_loss_list='train_loss'
val_loss_list='val_loss'

# training
def train(model, trainloader, optimizer, criterion):
    global logger
    global iter
    model.train()
    print('Training')
    # we will use this list to store the updated learning rates per epoch
    lrs=[]
    train_running_loss=0.0
    iters=len(trainloader)
    counter=0
    for i, data in tqdm(enumerate(trainloader), total=len(trainloader)):
        counter +=1
        image, labels=data
        image=image.cuda()
        labels=labels.cuda()
        optimizer.zero_grad()
        outputs=model(image)
        loss=criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        train_running_loss +=loss.item()
        logger.add_scalar("train_loss", loss.item(), iter)
        logger.add_scalar("lr", optimizer.param_groups[0]["lr"], iter)
        iter +=1

    epoch_loss=train_running_loss / counter
    return lrs, epoch_loss


# validation
def validate(model, testloader, criterion):
    model.eval()
    print('Validation')
    val_running_loss=0.0
    counter=0
    for i, data in tqdm(enumerate(testloader), total=len(testloader)):
        counter +=1

        image, labels=data
        image=image.cuda()
        labels=labels.cuda()
        outputs=model(image)
        loss=criterion(outputs, labels)
        val_running_loss +=loss.item()

    epoch_loss=val_running_loss / counter
    return epoch_loss


# start the training
train_loss, val_loss=[],[]
learning_rate_plot=[]
for epoch in range(epochs):
    print(f"[INFO]: Epoch{epoch+1}of{epochs}")
    lrs, train_epoch_loss=train(model, train_dataloader, optimizer,
                                  criterion)
    val_epoch_loss=validate(model, val_dataloader, criterion)
    logger.add_scalar("val_loss", val_epoch_loss, iter)
    train_loss.append(train_epoch_loss)
    val_loss.append(val_epoch_loss)
    learning_rate_plot.extend(lrs)
    print(f"Training loss:{train_epoch_loss:.3f}")
    print(f"Validation loss:{val_epoch_loss:.3f}")
    print('------------------------------------------------------------')