connect.py

import argparse
import numpy as np
import os
import sys
import tabulate
import torch
import torch.nn.functional as F

import data
import models
import utils


def main(args):
    os.makedirs(args.dir, exist_ok=True)
    with open(os.path.join(args.dir, 'chain.sh'), 'w') as f:
        f.write(' '.join(sys.argv))
        f.write('\n')

    torch.backends.cudnn.benchmark = True

    loaders, num_classes = data.loaders(
        args.dataset,
        args.data_path,
        args.batch_size,
        args.num_workers,
        args.transform,
        args.use_test,
        shuffle_train=False
    )

    architecture = getattr(models, args.model)
    base_model = architecture.base(num_classes, **architecture.kwargs)
    base_model.cuda()

    criterion = F.cross_entropy
    regularizer = utils.l2_regularizer(args.wd)

    def get_weights(model):
        return np.concatenate([p.data.cpu().numpy().ravel() for p in model.parameters()])

    T = (args.num_points - 1) * len(args.ckpt - 1) + 1
    ts = np.linspace(0.0, len(args.ckpt) - 1, T)
    tr_loss = np.zeros(T)
    tr_nll = np.zeros(T)
    tr_acc = np.zeros(T)
    te_loss = np.zeros(T)
    te_nll = np.zeros(T)
    te_acc = np.zeros(T)
    tr_err = np.zeros(T)
    te_err = np.zeros(T)

    columns = ['t', 'Train loss', 'Train nll', 'Train error (%)', 'Test nll', 'Test error (%)']

    alphas = np.linspace(0.0, 1.0, args.num_points)

    for path in args.ckpt:
        print(path)

    step = 0
    for i in range(len(args.ckpt) - 1):
        base_model.load_state_dict(torch.load(args.ckpt[i])['model_state'])
        w_1 = get_weights(base_model)

        base_model.load_state_dict(torch.load(args.ckpt[i + 1])['model_state'])
        w_2 = get_weights(base_model)
        for alpha in alphas[1 if i > 0 else 0:]:
            w = (1.0 - alpha) * w_1 + alpha * w_2
            offset = 0
            for parameter in base_model.parameters():
                size = np.prod(parameter.size())
                value = w[offset:offset + size].reshape(parameter.size())
                parameter.data.copy_(torch.from_numpy(value))
                offset += size

            utils.update_bn(loaders['train'], base_model)

            tr_res = utils.test(loaders['train'], base_model, criterion, regularizer)
            te_res = utils.test(loaders['test'], base_model, criterion, regularizer)

            tr_loss[step] = tr_res['loss']
            tr_nll[step] = tr_res['nll']
            tr_acc[step] = tr_res['accuracy']
            tr_err[step] = 100.0 - tr_acc[step]
            te_loss[step] = te_res['loss']
            te_nll[step] = te_res['nll']
            te_acc[step] = te_res['accuracy']
            te_err[step] = 100.0 - te_acc[step]

            values = [ts[step], tr_loss[step], tr_nll[step], tr_err[step], te_nll[step], te_err[step]]
            table = tabulate.tabulate([values], columns, tablefmt='simple', floatfmt='10.4f')
            if step % 40 == 0:
                table = table.split('\n')
                table = '\n'.join([table[1]] + table)
            else:
                table = table.split('\n')[2]
            print(table)
            step += 1

    np.savez(
        os.path.join(args.dir, 'chain.npz'),
        ts=ts,
        tr_loss=tr_loss,
        tr_nll=tr_nll,
        tr_acc=tr_acc,
        tr_err=tr_err,
        te_loss=te_loss,
        te_nll=te_nll,
        te_acc=te_acc,
        te_err=te_err,
    )


if __name__ == "__main":
    parser = argparse.ArgumentParser(description='Connect models with polychain')

    parser.add_argument('--dir', type=str, default='/tmp/chain/', metavar='DIR',
                        help='training directory (default: /tmp/chain)')

    parser.add_argument('--num_points', type=int, default=6, metavar='N',
                        help='number of points between models (default: 6)')

    parser.add_argument('--dataset', type=str, default='CIFAR10', metavar='DATASET',
                        help='dataset name (default: CIFAR10)')
    parser.add_argument('--use_test', action='store_true',
                        help='switches between validation and test set (default: validation)')
    parser.add_argument('--transform', type=str, default='VGG', metavar='TRANSFORM',
                        help='transform name (default: VGG)')
    parser.add_argument('--data_path', type=str, default=None, metavar='PATH',
                        help='path to datasets location (default: None)')
    parser.add_argument('--batch_size', type=int, default=128, metavar='N',
                        help='input batch size (default: 128)')
    parser.add_argument('--num_workers', type=int, default=4, metavar='N',
                        help='number of workers (default: 4)')

    parser.add_argument('--model', type=str, default=None, metavar='MODEL',
                        help='model name (default: None)')

    parser.add_argument('--ckpt', type=str, action='append', metavar='CKPT', required=True,
                        help='checkpoint to eval, pass all the models through this parameter')

    parser.add_argument('--wd', type=float, default=1e-4, metavar='WD',
                        help='weight decay (default: 1e-4)')

    args = parser.parse_args()
    main(args)