Function fully_checkpoint_sequential

PATH/core/utils.py:696–737

r"""Modified version of torch.utils.checkpoint.checkpoint_sequential for memory efficiency. It is assumed that at least one of the inputs have requires_grad=True, so we can checkpoint all of the segments at ease. Please refer to https://pytorch.org/docs/stable/checkpoint.html#torch.utils

(functions, segments, input, **kwargs)

    print('sync_print: rank {}, '.format(rank) + ' '.join(args), **kwargs)

def fully_checkpoint_sequential(functions, segments, input, **kwargs):
    r"""Modified version of torch.utils.checkpoint.checkpoint_sequential for memory efficiency.
    It is assumed that at least one of the inputs have requires_grad=True, so we can checkpoint
    all of the segments at ease.
    Please refer to https://pytorch.org/docs/stable/checkpoint.html#torch.utils.checkpoint.checkpoint_sequential
    for more details.

    -1 -> sqrt chunk checkpoint
    0  -> no checkpoint
    others ->
    """
    preserve = kwargs.pop('preserve_rng_state', True)
    if kwargs:
        raise ValueError("Unexpected keyword arguments: " + ",".join(arg for arg in kwargs))

    def run_function(start, end, functions):
        def forward(input):
            for j in range(start, end + 1):
                input = functions[j](input)
            return input
        return forward

    if isinstance(functions, torch.nn.Sequential):
        functions = list(functions.children())

    # no checkpoint
    if segments == 0:
        return run_function(0, len(functions) - 1, functions)(input)

    # auto determin the chunksize
    if segments < 0:
        segments = int(math.ceil(len(functions)))

    segments = min(segments, len(functions))
    segment_size = len(functions) // segments
    # the last chunk has to be non-volatile
    end = -1
    for start in range(0, segment_size * (segments - 1), segment_size):
        end = start + segment_size - 1
        input = checkpoint(run_function(start, end, functions), input)
#                           preserve_rng_state=preserve)
    return checkpoint(run_function(end + 1, len(functions) - 1, functions), input)#,
#                      preserve_rng_state=preserve)

def printlog(*args, **kwargs):