test_curve.py

import argparse
import torch

import curves
import data
import models

"""
This tests the max observed error between two endpoints
there is no apparent reason to use the curve module
this script is not important
"""
parser = argparse.ArgumentParser(description='Test DNN curve')

parser.add_argument('--dataset', type=str, default=None, 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', required=True,
                    help='model name (default: None)')

parser.add_argument('--curve', type=str, default=None, metavar='CURVE', required=True,
                    help='curve type to use (default: None)')
parser.add_argument('--num_bends', type=int, default=3, metavar='N',
                    help='number of curve bends (default: 3)')
parser.add_argument('--init_start', type=str, default=None, metavar='CKPT',
                    help='checkpoint to init start point (default: None)')
parser.add_argument('--init_end', type=str, default=None, metavar='CKPT',
                    help='checkpoint to init end point (default: None)')
parser.set_defaults(init_linear=True)
parser.add_argument('--init_linear_off', dest='init_linear', action='store_false',
                    help='turns off linear initialization of intermediate points (default: on)')

parser.add_argument('--seed', type=int, default=1, metavar='S', help='random seed (default: 1)')

args = parser.parse_args()

torch.backends.cudnn.benchmark = True
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)

if args.dataset is not None:
    loaders, num_classes = data.loaders(
        args.dataset,
        args.data_path,
        args.batch_size,
        args.num_workers,
        args.transform,
        args.use_test
    )
    loader = loaders['test']
else:
    num_classes = 10
    loader = [(torch.randn((args.batch_size, 3, 32, 32)), None) for i in range(20)]

architecture = getattr(models, args.model)

curve = getattr(curves, args.curve)
curve_model = curves.CurveNet(
    num_classes,
    curve,
    architecture.curve,
    args.num_bends,
    True,
    True,
    architecture_kwargs=architecture.kwargs,
)

base = [architecture.base(num_classes, **architecture.kwargs) for _ in range(2)]
for base_model, path, k in zip(base, [args.init_start, args.init_end], [0, args.num_bends - 1]):
    if path is not None:
        checkpoint = torch.load(path)
        print('Loading %s as point #%d' % (path, k))
        base_model.load_state_dict(checkpoint['model_state'])
    curve_model.import_base_parameters(base_model, k)

if args.init_linear:
    print('Linear initialization.')
    curve_model.init_linear()
curve_model = curve_model.cuda()
for base_model in base:
    base_model.cuda()

t = torch.FloatTensor([0.0]).cuda()
for base_model, t_value in zip(base, [0.0, 1.0]):
    print('T: %f' % t_value)
    t.data.fill_(t_value)
    curve_model.import_base_buffers(base_model)
    curve_model.eval()
    base_model.eval()

    max_error = 0.0
    for i, (input, _) in enumerate(loader):
        input = input.cuda(non_blocking=True)

        base_ouput = base_model(input)
        curve_output = curve_model(input, t)

        error = torch.max(torch.abs(base_ouput - curve_output)).item()
        print('Batch #%d. Error: %g' % (i, error))
        max_error = max(max_error, error)
    print('Max error: %g' % max_error)
    assert max_error < 1e-4, 'Error is too big (%g)' % max_error