|
| 1 | +# |
| 2 | +# Licensed to the Apache Software Foundation (ASF) under one |
| 3 | +# or more contributor license agreements. See the NOTICE file |
| 4 | +# distributed with this work for additional information |
| 5 | +# regarding copyright ownership. The ASF licenses this file |
| 6 | +# to you under the Apache License, Version 2.0 (the |
| 7 | +# "License"); you may not use this file except in compliance |
| 8 | +# with the License. You may obtain a copy of the License at |
| 9 | +# |
| 10 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | +# |
| 12 | +# Unless required by applicable law or agreed to in writing, |
| 13 | +# software distributed under the License is distributed on an |
| 14 | +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 15 | +# KIND, either express or implied. See the License for the |
| 16 | +# specific language governing permissions and limitations |
| 17 | +# under the License. |
| 18 | +# |
| 19 | + |
| 20 | + |
| 21 | +from singa import singa_wrap as singa |
| 22 | +from singa import opt |
| 23 | +from singa import tensor |
| 24 | +import argparse |
| 25 | +import train_cnn |
| 26 | + |
| 27 | +singa_dtype = {"float16": tensor.float16, "float32": tensor.float32} |
| 28 | + |
| 29 | +if __name__ == '__main__': |
| 30 | + # Use argparse to get command config: max_epoch, model, data, etc., for single gpu training |
| 31 | + parser = argparse.ArgumentParser( |
| 32 | + description='Training using the autograd and graph.') |
| 33 | + parser.add_argument('model', |
| 34 | + choices=['cnn', 'resnet', 'xceptionnet', 'mlp'], |
| 35 | + default='cnn') |
| 36 | + parser.add_argument('data', choices=['mnist', 'cifar10', 'cifar100'], default='mnist') |
| 37 | + parser.add_argument('-p', |
| 38 | + choices=['float32', 'float16'], |
| 39 | + default='float32', |
| 40 | + dest='precision') |
| 41 | + parser.add_argument('-m', |
| 42 | + '--max-epoch', |
| 43 | + default=10, |
| 44 | + type=int, |
| 45 | + help='maximum epochs', |
| 46 | + dest='max_epoch') |
| 47 | + parser.add_argument('-b', |
| 48 | + '--batch-size', |
| 49 | + default=64, |
| 50 | + type=int, |
| 51 | + help='batch size', |
| 52 | + dest='batch_size') |
| 53 | + parser.add_argument('-l', |
| 54 | + '--learning-rate', |
| 55 | + default=0.005, |
| 56 | + type=float, |
| 57 | + help='initial learning rate', |
| 58 | + dest='lr') |
| 59 | + parser.add_argument('-d', |
| 60 | + '--dist-option', |
| 61 | + default='plain', |
| 62 | + choices=['plain','half','partialUpdate','sparseTopK','sparseThreshold'], |
| 63 | + help='distibuted training options', |
| 64 | + dest='dist_option') # currently partialUpdate support graph=False only |
| 65 | + parser.add_argument('-s', |
| 66 | + '--sparsification', |
| 67 | + default='0.05', |
| 68 | + type=float, |
| 69 | + help='the sparsity parameter used for sparsification, between 0 to 1', |
| 70 | + dest='spars') |
| 71 | + parser.add_argument('-g', |
| 72 | + '--disable-graph', |
| 73 | + default='True', |
| 74 | + action='store_false', |
| 75 | + help='disable graph', |
| 76 | + dest='graph') |
| 77 | + parser.add_argument('-v', |
| 78 | + '--log-verbosity', |
| 79 | + default=0, |
| 80 | + type=int, |
| 81 | + help='logging verbosity', |
| 82 | + dest='verbosity') |
| 83 | + |
| 84 | + args = parser.parse_args() |
| 85 | + |
| 86 | + sgd = opt.SGD(lr=args.lr, momentum=0.9, weight_decay=1e-5, dtype=singa_dtype[args.precision]) |
| 87 | + sgd = opt.DistOpt(sgd) |
| 88 | + |
| 89 | + train_cnn.run(sgd.global_rank, sgd.world_size, sgd.local_rank, args.max_epoch, |
| 90 | + args.batch_size, args.model, args.data, sgd, args.graph, |
| 91 | + args.verbosity, args.dist_option, args.spars, args.precision) |
0 commit comments