train_multigpu.py

import os
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'
os.environ['JAX_PLATFORM_NAME'] = 'gpu'
# os.environ['XLA_PYTHON_CLIENT_ALLOCATOR'] = 'platform'
# os.environ['NVIDIA_TF32_OVERRIDE'] = '1'


import jax
import jax.lax
# from jax import random, grad, jit, vmap, value_and_grad
import numpy as np
import jax.numpy as jnp
import functools

import torch
from torchvision import datasets, transforms
from tqdm import trange, tqdm
import PIL.Image

import haiku as hk
import optax
import jmp

from typing import Any, Iterable, Mapping, NamedTuple, Tuple
import atexit
import resnet_cifar
import tree
import shutil
import einops


class FLAGS(NamedTuple):
    KEY = jax.random.PRNGKey(1)
    BATCH_SIZE = 128
    DATA_ROOT = '/workspace/data/'
    LOG_ROOT = '/workspace/runs/'
    MAX_EPOCH = 200
    INIT_LR = 1e-1
    N_WORKERS = 4
    MNIST_MEAN = (0.1307,)
    MNIST_STD = (0.3081,)
    CIFAR10_MEAN = (0.4914, 0.4822, 0.4465)
    CIFAR10_STD = (0.2023, 0.1994, 0.2010)
    CIFAR100_MEAN = (0.5071, 0.4867, 0.4408)
    CIFAR100_STD = (0.2675, 0.2565, 0.2761)
    IMAGENET_MEAN = (0.485, 0.456, 0.406)
    IMAGENET_STD = (0.229, 0.224, 0.225)
    NET = resnet_cifar.ResNet32


# shutil.copytree('./', FLAGS.LOG_ROOT, dirs_exist_ok=True)


def tprint(obj):
    tqdm.write(obj.__str__())


class TrainState(NamedTuple):
    params: hk.Params
    state: hk.State
    opt_state: optax.OptState
    loss_scale: jmp.LossScale


class MultiEpochsDataLoader(torch.utils.data.DataLoader):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._DataLoader__initialized = False
        self.batch_sampler = _RepeatSampler(self.batch_sampler)
        self._DataLoader__initialized = True
        self.iterator = super().__iter__()

    def __len__(self):
        return len(self.batch_sampler.sampler)

    def __iter__(self):
        for i in range(len(self)):
            yield next(self.iterator)


class _RepeatSampler(object):
    """ Sampler that repeats forever.
    Args:
        sampler (Sampler)
    """

    def __init__(self, sampler):
        self.sampler = sampler

    def __iter__(self):
        while True:
            yield from iter(self.sampler)


def numpy_collate(batch):
    if isinstance(batch[0], np.ndarray):
        return np.stack(batch)
    elif isinstance(batch[0], (tuple, list)):
        transposed = zip(*batch)
        return [numpy_collate(samples) for samples in transposed]
    else:
        return np.array(batch)


class ArrayNormalize(torch.nn.Module):
    def __init__(self, mean, std):
        super().__init__()
        self.mean = mean
        self.std = std

    def __call__(self, arr: np.ndarray) -> np.ndarray:
        assert isinstance(arr, np.ndarray), f'Input should be ndarray. Got {type(arr)}.'
        assert arr.ndim >= 3, f'Expected array to be a image of size (..., H, W, C). Got {arr.shape}.'

        dtype = arr.dtype
        mean = np.asarray(self.mean, dtype=dtype)
        std = np.asarray(self.std, dtype=dtype)
        if (std == 0).any():
            raise ValueError(
                f'std evaluated to zero after conversion to {dtype}, leading to division by zero.')
        if mean.ndim == 1:
            mean = mean.reshape(1, 1, -1)
        if std.ndim == 1:
            std = std.reshape(1, 1, -1)
        arr -= mean
        arr /= std
        return arr


class ToArray(torch.nn.Module):
    '''convert image to float and 0-1 range'''
    dtype = np.float32

    def __call__(self, x):
        assert isinstance(x, PIL.Image.Image)
        x = np.asarray(x, dtype=self.dtype)
        x /= 255.0
        return x


def softmax_cross_entropy(logits, labels):
    logp = jax.nn.log_softmax(logits)
    loss = -jnp.take_along_axis(logp, labels[:, None], axis=-1)
    return loss


def l2_loss(params):
    # l2_params = jax.tree_util.tree_leaves(params)
    l2_params = [p for ((mod_name, _), p) in tree.flatten_with_path(
        params) if 'batchnorm' not in mod_name]
    return 0.5 * sum(jnp.sum(jnp.square(p)) for p in l2_params)


def forward(images, is_training: bool):
    net = FLAGS.NET(num_classes=10, bn_config={'decay_rate': 0.9})
    return net(images, is_training=is_training)


@jax.jit
def ema_update(params, params_avg):
    '''polyak averaging'''
    params_avg = optax.incremental_update(params, params_avg, step_size=0.001)
    return params_avg


def main():

    normalize = ArrayNormalize(FLAGS.CIFAR10_MEAN, FLAGS.CIFAR10_STD)
    transform_train = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        ToArray(),
        normalize,
    ])
    transform_test = transforms.Compose([
        ToArray(),
        normalize,
    ])
    train_dataset = datasets.CIFAR10(
        FLAGS.DATA_ROOT, train=True, download=True, transform=transform_train)
    test_dataset = datasets.CIFAR10(
        FLAGS.DATA_ROOT, train=False, transform=transform_test)
    train_loader = MultiEpochsDataLoader(
        train_dataset,
        batch_size=FLAGS.BATCH_SIZE,
        shuffle=True,
        drop_last=True,
        num_workers=FLAGS.N_WORKERS,
        collate_fn=numpy_collate,
    )
    test_loader = MultiEpochsDataLoader(
        test_dataset,
        batch_size=FLAGS.BATCH_SIZE*2,
        shuffle=False,
        drop_last=False,
        num_workers=FLAGS.N_WORKERS,
        collate_fn=numpy_collate,
    )

    ## INITIALIZE MIXED PRECISION ##
    mp_policy = jmp.Policy(param_dtype = jnp.float32,
                            compute_dtype = jnp.float16,
                            output_dtype = jnp.float32)
    mp_bn_policy = jmp.Policy(param_dtype = jnp.float32,
                            compute_dtype = jnp.float32,
                            output_dtype = jnp.float16)
    loss_scale = jmp.DynamicLossScale(2.0**15)
    hk.mixed_precision.set_policy(FLAGS.NET, mp_policy)
    hk.mixed_precision.set_policy(hk.BatchNorm, mp_bn_policy)


    ## INITIALIZE MODEL ##
    model = hk.transform_with_state(forward)
    # model = hk.without_apply_rng(model)

    sample_input = jnp.ones((1, 32, 32, 3))
    params, state = model.init(FLAGS.KEY, sample_input, is_training=True)
    tprint(
        sum([p.size for p in jax.tree_util.tree_leaves(params)])
    )

    ## INITIALIZE OPTIMIZER ##
    learning_rate_fn = optax.cosine_onecycle_schedule(
        transition_steps=len(train_loader) * FLAGS.MAX_EPOCH,
        peak_value=FLAGS.INIT_LR,
        pct_start=0.3,
        div_factor=25.0,
        final_div_factor=1e4
    )
    # learning_rate_fn = optax.cosine_decay_schedule(
    #     init_value=FLAGS.INIT_LR,
    #     decay_steps=len(train_loader) * FLAGS.MAX_EPOCH,
    #     alpha=0.0
    #     )
    # learning_rate_fn = optax.piecewise_constant_schedule(
    #     init_value=FLAGS.INIT_LR,
    #     boundaries_and_scales={
    #         len(train_loader) * 100: 0.1,
    #         len(train_loader) * 150: 0.1,
    #     }
    # )
    optimizer = optax.sgd(learning_rate_fn, momentum=0.9, nesterov=False)
    opt_state = optimizer.init(params)

    if jax.default_backend() == 'gpu':
        n_devices = jax.local_device_count()
    else:
        n_devices = 1

    ## INITIALIZE TRAIN STATE ##
    train_state = TrainState(params, state, opt_state, loss_scale)
    train_state = jax.device_put_replicated(train_state, jax.local_devices())


    @functools.partial(jax.pmap, axis_name='i', donate_argnums=(0,))
    def train_step(train_state: TrainState, batch: dict):
        params, state, opt_state, loss_scale = train_state
        input, target = batch['image'], batch['label']
        def loss_fn(p):
            logits, state_new = model.apply(
                p, state, FLAGS.KEY, input, is_training=True)
            ce_loss = softmax_cross_entropy(logits, target).mean()
            loss = ce_loss + 1e-4 * l2_loss(p)
            return loss_scale.scale(loss), (loss, state_new)
        grads, (val, state) = jax.grad(loss_fn, has_aux=True)(params)
        grads = jax.lax.pmean(grads, axis_name='i')
        grads = loss_scale.unscale(grads)
        deltas, opt_state = optimizer.update(grads, opt_state)
        params = optax.apply_updates(params, deltas)

        # adjust loss scale
        grads_finite = jmp.all_finite(grads)
        loss_scale = loss_scale.adjust(grads_finite)
        params, state, opt_state = jmp.select_tree(
            grads_finite,
            (params, state, opt_state),
            train_state[:-1])

        train_state = TrainState(params, state, opt_state, loss_scale)
        return train_state, val


    @jax.jit
    def eval_step(train_state: TrainState, batch: dict):
        params, state, _, _ = jax.tree_map(lambda x: x[0], train_state)
        input, target = batch['image'], batch['label']
        logits, _ = model.apply(params, state, FLAGS.KEY, input, is_training=False)
        correct = jnp.argmax(logits, -1) == target
        loss = softmax_cross_entropy(logits, target)
        return correct, loss


    def evaluate(dataloader):
        corrects = []
        losses = []
        for input, target in dataloader:
            batch = {
                'image': input,
                'label': target,
            }
            correct, loss = eval_step(train_state, batch)
            corrects.append(correct)
            losses.append(loss)
        avg_acc = jnp.concatenate(corrects).mean()
        avg_loss = jnp.concatenate(losses).mean()
        return avg_acc, avg_loss


    def save_pickle():
        state_dict = {
            'params': train_state.params,
            'state' : train_state.state,
            'opt_state': train_state.opt_state,
            'loss_scale': train_state.loss_scale,
        }
        pickle_path = os.path.join(FLAGS.LOG_ROOT, 'model.pickle')
        torch.save(state_dict, pickle_path)
        tprint(f'[SAVE] {pickle_path}')


    atexit.register(save_pickle)

    for epoch in trange(FLAGS.MAX_EPOCH, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}', smoothing=1.):
        for input, target in train_loader:
            batch = {
                'image': einops.rearrange(input, '(n b) h w c -> n b h w c', n=n_devices),
                'label': einops.rearrange(target, '(n b) -> n b', n=n_devices),
            }
            train_state, train_loss = train_step(train_state, batch)
        acc, loss = evaluate(test_loader)
        last_lr = learning_rate_fn(jax.tree_map(lambda x: x[0], train_state).opt_state[-1].count).item()
        tprint(f'[{epoch}/{FLAGS.MAX_EPOCH}] LR: {last_lr:.3f} | Train Loss {train_loss.mean():.3f} | Test Loss {loss:.3f} Acc: {acc:.3f}')

        if epoch % 20 == 0:
            save_pickle()


if __name__ == '__main__':
    main()