torch-nebula 0.1.0rc5
Save and load training snapshot
License: MIT
Requires Python: >=3.6

## Description

Save and load PyTorch checkpoint faster.
It saves checkpoint to shared memory and return immediately.
The replication and persisting will be done asynchronously by background service.

## Install
```
python -m pip install --upgrade pip
pip install torch-nebula
```

## API
```
def init(config=NebulaConfig()):
    # launch nebula services
    # it MUST be called before calling other nebula functions
    # config:
    #   configurations such as Tier2 storage path
    #   set to default if not given
    pass

class NebulaConfig(object):
    def __init__(self, replica_dir=None):
        # replica_dir:
        #   Path of Tier2 storage
        pass

class Checkpoint(object):
    def __init__(self, tag: str, num_of_files: int = -1, client: NebulaClient = None):
        # tag: 
        #   an unique name for a Checkpoint
        #   tag usually will be format of "global_step1200"
        # num_of_files:
        #   total number of files in this Checkpoint
        #   the number of calling "save" on this Checkpoint MUST be equal to this
        pass
        
    def save(self, name, state):
        # save state to cache under /tag/name.
        pass
    
    def load(self, name, map_location=None):
        # get state under /tag/name from agent and return it.
        pass

def get_latest_checkpoint():
    # return latest VALID Checkpoint.
    # otherwise, return None
    pass

def get_checkpoint(tag: str = None):
    # return Checkpoint by tag, if snapshot exist
    # no matter it's valid or not
    # otherwise, return None
    pass

def list_checkpoints(tag: str = None):
    # list checkpoints information in json format
    # if tag is set, only list that checkpoint.
```

#**Example usage**
The below code could be launced in 2 nodes. Each node uses only 1 GPU.

Launching 2 processes from scratch:
```
On node-0:
    CUDA_VISIBLE_DEVICES=0 python resnet50_example_ddp.py 0 2 save
On node-1:
    CUDA_VISIBLE_DEVICES=0 python resnet50_example_ddp.py 1 2 save
```
Launching 2 processes from loading latest checkpoint:
```
On node-0:
    CUDA_VISIBLE_DEVICES=0 python resnet50_example_ddp.py 0 2 load
On node-1:
    CUDA_VISIBLE_DEVICES=0 python resnet50_example_ddp.py 1 2 load
```
Code of resnet50_example_ddp.py:
```
import os
import sys

import torch
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim
import torch.utils.data
import torchvision
import torchvision.models as models
import torchvision.transforms as T
from torch.nn.parallel import DistributedDataParallel as DDP
import torch_nebula as tn


def example(rank, world_size, is_load=False):
    tn.init()  # MUST be called before using nebula.

    transform = T.Compose([T.Resize(256), T.CenterCrop(224), T.ToTensor()])
    trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                            download=True, transform=transform)
    train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=32, sampler=train_sampler,
                                              shuffle=False, num_workers=4)

    gpu_id = 0  # In this demo, only one GPU is used per node.
    model = models.resnet50(pretrained=False)
    if is_load:
        ckpt = tn.get_latest_checkpoint()
        if ckpt is None:
            print("No available checkpoint!")
            exit(-1)
        states_model = ckpt.load("rank{}_model".format(rank),
                                 map_location=torch.device('cuda:0'))
        model.load_state_dict(states_model["model"])
    model = model.to(gpu_id)
    model.cuda()
    cudnn.benchmark = True
    print("before DDP")
    model = DDP(model, device_ids=[gpu_id])
    print("after DDP")
    criterion = nn.CrossEntropyLoss().to(gpu_id)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
    if is_load:
        states_optimizer = ckpt.load("rank{}_optimizer".format(rank),
                                     map_location=torch.device('cpu'))
        optimizer.load_state_dict(states_optimizer["optimizer"])
    model.train()

    start_epoch = 0
    start_step = 0
    global_step = 0
    if is_load:
        print("loading from Checkpoint with tag={}".format(ckpt.tag))
        start_epoch = states_model["epoch"]
        start_step = states_model["step"]
        global_step = states_model["global_step"]

    for epoch in range(start_epoch, 100000):
        for step, data in enumerate(trainloader, start_step):
            inputs, labels = data[0].to(gpu_id), data[1].to(gpu_id)
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            print("rank{} step{} loss={}".format(rank, global_step, loss.item()))

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if global_step % 100 == 0:
                print("rank {}: Saving model...".format(rank))
                tag = "global_step{}".format(global_step)
                # Each process saves 2 files: model and optimizer.
                ckpt = tn.Checkpoint(tag=tag, num_of_files=world_size * 2)
                m = model.module  # DDP needs this.
                ckpt.save("rank{}_model".format(rank), {
                    "model": m.state_dict(),
                    "epoch": epoch,
                    "step": step,
                    "global_step": global_step
                })
                ckpt.save("rank{}_optimizer".format(rank), {
                    "optimizer": optimizer.state_dict()
                })

            global_step += 1


def init_process(rank, size, is_load, fn, backend='nccl'):
    dist.init_process_group(backend,
                            init_method="tcp://{}:{}".format('192.168.0.28', '29509'),
                            rank=rank,
                            world_size=size)
    fn(rank, size, is_load=is_load)


if __name__=="__main__":
    rank = int(sys.argv[1])
    size = int(sys.argv[2])
    is_load = (sys.argv[3] == "load")
    print("rank={}, size={}".format(rank, size))

    mp.set_start_method("spawn")
    p = mp.Process(target=init_process, args=(rank, size, is_load, example))
    p.start()
    p.join()

```

##**V0.1.0**
- support save and get latest version