FM / sync

/usr/local/lib64/python3.6/site-packages/torch/distributed/pipeline/sync/

/usr/local/lib64/python3.6/site-packages/torch/distributed/pipeline/sync

Name	Size	Mode	Actions
skip/	-	0755	rm
_balance/	-	0755	rm
__pycache__/	-	0755	rm
batchnorm.py	5603	0644	edit dl rm
checkpoint.py	11464	0644	edit dl rm
copy.py	3716	0644	edit dl rm
dependency.py	1694	0644	edit dl rm
microbatch.py	7427	0644	edit dl rm
phony.py	1536	0644	edit dl rm
pipe.py	17883	0644	edit dl rm
pipeline.py	9638	0644	edit dl rm
stream.py	3807	0644	edit dl rm
utils.py	1147	0644	edit dl rm
worker.py	4284	0644	edit dl rm
__init__.py	451	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/distributed/pipeline/sync/utils.py (1147B)

from torch import nn
from typing import List

def partition_model(
        module: nn.Sequential,
        balance: List[int],
        devices: List[int] = None):
    """
    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)

from torch import nn from typing import List def partition_model( module: nn.Sequential, balance: List[int], devices: List[int] = None): """ Given an :class:`nn.Sequential ` module, partitions the model across multiple GPU devices according the provided ``balance`` and ``devices``. Args: module (:class:`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)