コード例 #1
0
def generic_main(args):
    import sys
    from MixedPrecision.tools.utils import summary
    import MixedPrecision.tools.loaders as loaders

    sys.stderr = sys.stdout

    torch.set_num_threads(args.workers)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    utils.set_use_gpu(args.gpu, not args.no_bench_mode)
    utils.set_use_half(args.half)
    utils.setup(args)
    utils.show_args(args)

    model = models.__dict__[args.model]()

    summary(model, input_size=(3, 224, 224), batch_size=args.batch_size)

    data = loaders.load_dataset(args, train=True)

    if args.warmup:
        train(args, model, data, args.model, is_warmup=True)

    train(args, model, data, args.model, is_warmup=False)

    sys.exit(0)
コード例 #2
0
def main():
    import sys
    from MixedPrecision.pytorch.mnist_fully_connected import load_mnist
    from MixedPrecision.pytorch.mnist_fully_connected import train
    from MixedPrecision.pytorch.mnist_fully_connected import init_weights
    from MixedPrecision.tools.args import get_parser
    from MixedPrecision.tools.utils import summary
    import MixedPrecision.tools.utils as utils

    torch.manual_seed(0)
    torch.cuda.manual_seed_all(0)

    parser = get_parser()
    args = parser.parse_args()

    utils.set_use_gpu(args.gpu)
    utils.set_use_half(args.half)

    shape = (1, 28, 28)
    if args.fake:
        shape = args.shape

    for k, v in vars(args).items():
        print('{:>30}: {}'.format(k, v))

    try:
        current_device = torch.cuda.current_device()
        print('{:>30}: {}'.format('GPU Count', torch.cuda.device_count()))
        print('{:>30}: {}'.format('GPU Name',
                                  torch.cuda.get_device_name(current_device)))
    except:
        pass

    model = MnistConvolution(input_shape=shape,
                             conv_num=args.conv_num,
                             kernel_size=args.kernel_size,
                             explicit_permute=args.permute)

    model.float()
    model.apply(init_weights)
    model = utils.enable_cuda(model)
    summary(model, input_size=(shape[0], shape[1], shape[2]))
    model = utils.enable_half(model)

    train(
        args, model,
        load_mnist(args,
                   hwc_permute=args.permute,
                   fake_data=args.fake,
                   shape=shape))

    sys.exit(0)
コード例 #3
0
ファイル: exp1b.py プロジェクト: Delaunay/MixedPrecision
    from MixedPrecision.tools.optimizer import OptimizerAdapter
    import MixedPrecision.tools.utils as utils

    from apex.fp16_utils import network_to_half

    sys.stderr = sys.stdout

    parser = get_parser()
    args = parser.parse_args()

    torch.set_num_threads(args.workers)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    utils.set_use_gpu(args.gpu, not args.no_bench_mode)
    utils.set_use_half(args.half)
    utils.show_args(args)

    data_loader = load_dataset(args, train=True)

    model = utils.enable_cuda(HybridClassifier())

    if args.half:
        model = network_to_half(model)

    criterion = utils.enable_cuda(HybridLoss())

    optimizer = torch.optim.SGD(
        model.parameters(),
        args.lr,
        momentum=args.momentum,
コード例 #4
0
def main():
    # This does not work but this is what the documentation says to do...
    #try:
    #    import torch.multiprocessing as multiprocessing
    #    multiprocessing.set_start_method('spawn')
    #except Exception as e:
    #    print(e)

    import MixedPrecision.tools.utils as utils
    import argparse

    parser = argparse.ArgumentParser(description='Data loader Benchmark')

    parser.add_argument('--data',
                        type=str,
                        metavar='DIR',
                        help='path to the dataset location')

    parser.add_argument('-j',
                        '--workers',
                        default=4,
                        type=int,
                        metavar='N',
                        help='number of data loading workers (default: 4)')

    parser.add_argument('--epochs',
                        default=10,
                        type=int,
                        metavar='N',
                        help='number of total epochs to run')

    parser.add_argument('-b',
                        '--batch-size',
                        default=256,
                        type=int,
                        metavar='N',
                        help='mini-batch size (default: 256)')

    parser.add_argument('--prof',
                        dest='prof',
                        type=int,
                        default=10,
                        help='Only run N iterations for profiling.')

    parser.add_argument(
        '--loader',
        type=str,
        default='pytorch',
        help='The kind of loader to use (torch, prefetch, benzina, dali, zip)')

    parser.add_argument('--async',
                        action='store_true',
                        default=False,
                        help='Use AsyncPrefetcher')

    args = parser.parse_args()

    utils.set_use_gpu(True, True)
    utils.set_use_half(True)

    utils.show_args(args)

    benchmark_loader(args)