File size: 1,290 Bytes
c61ccee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
from torch import nn
from typing import List, Optional

__all__ = ["partition_model"]

def partition_model(

        module: nn.Sequential,

        balance: List[int],

        devices: Optional[List[int]] = None):
    """

    Partions the model accross multiple GPU devices.



    Given an :class:`nn.Sequential <torch.nn.Sequential>` module, partitions

    the model across multiple GPU devices according the provided ``balance``

    and ``devices``.



    Args:

        module (:class:`nn.Sequential <torch.nn.Sequential>`):

            Sequential model representing the pipe.

        balance (List[int]):

            List indicating the number of layers in each partition.

        devices (List[int], optional):

            List indicating the device to use for each partition. Defaults to

            ``range(len(balance))``

    """
    device_idx = 0
    pipe_idx = 0
    balanced_pipe = []
    for num_layers in balance:
        layers = []
        for i in range(num_layers):
            layers.append(module[pipe_idx])
            pipe_idx += 1
        device = device_idx if devices is None else devices[device_idx]
        balanced_pipe.append(nn.Sequential(*layers).to(device))
        device_idx += 1

    return nn.Sequential(*balanced_pipe)