Fastai Distributed Language Model

dist-imdb.py

from fastai import *
from fastai.text import *

from fastai.distributed import *
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int)
args = parser.parse_args()
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')

path = Path("/root/.fastai/data/imdb")

def random_seed(seed_value, use_cuda):
    np.random.seed(seed_value) # cpu vars
    torch.manual_seed(seed_value) # cpu  vars
    random.seed(seed_value) # Python
    if use_cuda:
        torch.cuda.manual_seed(seed_value)
        torch.cuda.manual_seed_all(seed_value) # gpu vars
        torch.backends.cudnn.deterministic = True  #needed
        torch.backends.cudnn.benchmark = False

random_seed(42, True)

data_lm = load_data(path, 'data_lm.pkl', bs=80)
learn = language_model_learner(data_lm, AWD_LSTM, drop_mult=0.3)
learn = learn.to_distributed(args.local_rank)

learn.fit_one_cycle(1, 1e-2)
learn.unfreeze()
learn.fit_one_cycle(10, 1e-3, moms=(0.8, 0.7))