Mac Pro M1测试PyTorch GPU

3,355 阅读4分钟

携手创作,共同成长!这是我参与「掘金日新计划 · 8 月更文挑战」的第1天,点击查看活动详情

前言

这几天暑假回家被社区集中隔离,进来就带了台笔记本每天实在是太无聊了。想起来之前刷到最新版本的Pytorch貌似已经支持M1芯片的GPU加速,趁着有时间动手试试看效果怎么样(目前还是预览版所以对性能啥的也没有期待纯粹好玩

安装

安装部分还是挺简单的,之前装过tf的M1适配环境,所以这次继续在之前这个环境装个Pytorch。

第一步还是打开官网Pytorch官网,找到安装部分选择好平台啥的,最重要的是选择预览版本。截图如下

pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu 运行这行命令进行安装,然后下面来开始测试一下

开始测试

首先查看相关的信息

import torch
torch.__version__

print(f"{torch.backends.mps.is_available(),torch.backends.mps.is_built()}")

a=torch.rand(5).to("mps")
a

在这里插入图片描述

可以看到Pytorch为M1系列添加了新的后端,不是cuda而是mps,如果去看官方文档还会发现很多其他接口。目前就这个效果来看确实能使用M1的GPU,但是效果如何呢?我本来想类似于cuda一样使用torch.cuda.list_gpu_processes()来查看GPU相关信息,但是貌似他们还没有写mps对应的方法。那就写个简单的demo测试一下,不知道M1的GPU到底适配性如何了。

这里先贴一张系统信息的图

在这里插入图片描述

内存一共16G,不知道拉满能用多少

这次的测试代码就是利用resnet18微调训练cifar10,优化器还是使用Ranger,代码如下

def get_dataloader(batch_size):
    data_transform = {
        "train": transforms.Compose([transforms.Resize(96),
                                     transforms.ToTensor()]),
        "val": transforms.Compose([transforms.Resize(96),
                                   transforms.ToTensor()])
    }
    train_dataset = torchvision.datasets.CIFAR10('./p10_dataset', train=True, transform=data_transform["train"], download=True)
    test_dataset = torchvision.datasets.CIFAR10('./p10_dataset', train=False, transform=data_transform["val"], download=True)
    print('训练数据集长度: {}'.format(len(train_dataset)))
    print('测试数据集长度: {}'.format(len(test_dataset)))
    # DataLoader创建数据集
    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)
    return train_dataloader,test_dataloader

def show_pic(dataloader):#展示dataloader里的6张图片
    examples = enumerate(dataloader)  # 组合成一个索引序列
    batch_idx, (example_data, example_targets) = next(examples)
    classes = ('airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    fig = plt.figure()
    for i in range(6):
        plt.subplot(2, 3, i + 1)
        # plt.tight_layout()
        img = example_data[i]
        print('pic shape:',img.shape)
        img = img.swapaxes(0, 1)
        img = img.swapaxes(1, 2)
        plt.imshow(img, interpolation='none')
        plt.title(classes[example_targets[i].item()])
        plt.xticks([])
        plt.yticks([])
    plt.show()

def get_net(): #获得预训练模型并冻住前面层的参数
    net = timm.create_model('resnet18', pretrained=True, num_classes=10)
    #print(summary(net, input_size=(3, 224, 224)))
    '''Freeze all layers except the last layer(fc or classifier)'''
    for param in net.parameters():
        param.requires_grad = False
    # nn.init.xavier_normal_(model.fc.weight)
    # nn.init.zeros_(model.fc.bias)
    net.fc.weight.requires_grad = True
    net.fc.bias.requires_grad = True
    return net

def train(net, loss, train_dataloader, valid_dataloader, device, batch_size, num_epoch, lr, lr_min, optim='sgd', init=True, scheduler_type='Cosine'):
    def init_xavier(m):
        #if type(m) == nn.Linear or type(m) == nn.Conv2d:
        if type(m) == nn.Linear:
            nn.init.xavier_normal_(m.weight)

    if init:
        net.apply(init_xavier)

    print('training on:', device)
    net.to(device)

    if optim == 'sgd':
        optimizer = torch.optim.SGD((param for param in net.parameters() if param.requires_grad), lr=lr,
                                    weight_decay=0)
    elif optim == 'adam':
        optimizer = torch.optim.Adam((param for param in net.parameters() if param.requires_grad), lr=lr,
                                     weight_decay=0)
    elif optim == 'adamW':
        optimizer = torch.optim.AdamW((param for param in net.parameters() if param.requires_grad), lr=lr,
                                      weight_decay=0)
    elif optim == 'ranger':
        optimizer = Ranger21((param for param in net.parameters() if param.requires_grad), lr=lr,weight_decay=0,num_epochs=num_epoch,num_batches_per_epoch=len(train_dataloader),use_warmup=False,use_madgrad=False)
    if scheduler_type == 'Cosine':
        scheduler = CosineAnnealingLR(optimizer, T_max=num_epoch, eta_min=lr_min)
    
    if scheduler_type == 'Cyclic':
        scheduler =CyclicLR(optimizer,base_lr=lr_min,max_lr=lr)

    train_losses = []
    train_acces = []
    eval_acces = []
    best_acc = 0.0
    for epoch in range(num_epoch):

        print("——————第 {} 轮训练开始——————".format(epoch + 1))

        # 训练开始
        net.train()
        train_acc = 0
        for batch in tqdm(train_dataloader):
            imgs, targets = batch
            imgs = imgs.to(device)
            targets = targets.to(device)
            output = net(imgs)

            Loss = loss(output, targets)
        
            optimizer.zero_grad()
            Loss.backward()
            optimizer.step()

            _, pred = output.max(1)
            num_correct = (pred == targets).sum().item()
            acc = num_correct / (batch_size)
            train_acc += acc
        scheduler.step()
        print("epoch: {}, Loss: {}, Acc: {}".format(epoch+1, Loss.item(), train_acc / len(train_dataloader)))
        train_acces.append(train_acc / len(train_dataloader))
        train_losses.append(Loss.item())

        # 测试步骤开始
        net.eval()
        eval_loss = 0
        eval_acc = 0
        with torch.no_grad():
            for imgs, targets in valid_dataloader:
                imgs = imgs.to(device)
                targets = targets.to(device)
                output = net(imgs)
                Loss = loss(output, targets)
                _, pred = output.max(1)
                num_correct = (pred == targets).sum().item()
                eval_loss += Loss
                acc = num_correct / imgs.shape[0]
                eval_acc += acc

            eval_losses = eval_loss / (len(valid_dataloader))
            eval_acc = eval_acc / (len(valid_dataloader))
            if eval_acc > best_acc:
                best_acc = eval_acc
                torch.save(net.state_dict(),'best_acc.pth')
            eval_acces.append(eval_acc)
            print("整体验证集上的Loss: {}".format(eval_losses))
            print("整体验证集上的正确率: {}".format(eval_acc))
    return train_losses, train_acces, eval_acces

def show_acces(train_losses, train_acces, valid_acces, num_epoch):#对准确率和loss画图显得直观
    plt.plot(1 + np.arange(len(train_losses)), train_losses, linewidth=1.5, linestyle='dashed', label='train_losses')
    plt.plot(1 + np.arange(len(train_acces)), train_acces, linewidth=1.5, linestyle='dashed', label='train_acces')
    plt.plot(1 + np.arange(len(valid_acces)), valid_acces, linewidth=1.5, linestyle='dashed', label='valid_acces')
    plt.grid()
    plt.xlabel('epoch')
    plt.xticks(range(1, 1 + num_epoch, 1))
    plt.legend()
    plt.show()

if __name__ == '__main__':
    train_dataloader, test_dataloader = get_dataloader(batch_size=1024)
    show_pic(train_dataloader)
    device = torch.device("mps")
    net = get_net()
    loss = nn.CrossEntropyLoss()
    train_losses, train_acces, eval_acces = train(net, loss, train_dataloader, test_dataloader, device, batch_size=1024, num_epoch=10, lr=0.1, lr_min=1e-5, optim='ranger',scheduler_type='Cyclic',init=False)
    show_acces(train_losses, train_acces, eval_acces, num_epoch=10)

最终结果 在这里插入图片描述

当然,我还拿出了很多年前本科大二买的一台1060笔记本一起跑这段代码,结果如下 在这里插入图片描述

老的笔记本显存只有6G,但是从结果来看由于图片很小所以显存没占满并且速度还比M1快很多,差不多只用M123\frac{2}{3}的时间就可以训练结束,这个差距还是挺大的。

总结

尝试了一次训练后,其实对M1挺失望的。或许它的优势也就是显存大一点而已,但是目前新款的笔记本或者游戏本显存也有12G,并且价格还没有它高。我们一般在笔记本上搞深度学习只是为了调通代码测试一下大概的效果,真正训练还是放到服务器上跑。所以唯一值得欣慰一点的就是用M1会减轻一点身体负担,毕竟通常游戏本加电源啥的那一套相比Mac来说还是太重了,其他的对于深度学习来说实在想不出优点(如果有人爱面子啥的那勉强也能凑一个)。

综上,即使现在主流的两大框架TensorFlow2Pytorch都支持M1及以上芯片的GPU加速,但是如果真的为了深度学习来专门买那就是智商税。平时搞搞前后端开发啥的很推荐,易携带而且续航,无噪音,编译速度快;搞深度学习的话不如花一半左右的钱买个高档显存大点的游戏本,效果会好很多。