• Table of Contents

Description

• VRAM 不夠時可以用。

Basic

只需要將模型的每一層分開存放在不同的 GPU 上

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

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = torch.nn.Linear(10, 10).to('cuda:0')
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(10, 5).to('cuda:1')

    def forward(self, x):
        x = self.relu(self.net1(x.to('cuda:0')))
        return self.net2(x.to('cuda:1'))

Reference

• Single-Machine Model Parallel Best Practices — PyTorch Tutorials 1.12.1+cu102 documentation