()
| 86 | |
| 87 | |
| 88 | def get_config(): |
| 89 | nr_tower = max(get_num_gpu(), 1) |
| 90 | batch = args.batch |
| 91 | total_batch = batch * nr_tower |
| 92 | if total_batch != 128: |
| 93 | logger.warn("AlexNet needs to be trained with a total batch size of 128.") |
| 94 | BASE_LR = 0.01 * (total_batch / 128.) |
| 95 | |
| 96 | logger.info("Running on {} towers. Batch size per tower: {}".format(nr_tower, batch)) |
| 97 | dataset_train = get_data('train', batch) |
| 98 | dataset_val = get_data('val', batch) |
| 99 | |
| 100 | infs = [ClassificationError('wrong-top1', 'val-error-top1'), |
| 101 | ClassificationError('wrong-top5', 'val-error-top5')] |
| 102 | callbacks = [ |
| 103 | ModelSaver(), |
| 104 | GPUUtilizationTracker(), |
| 105 | EstimatedTimeLeft(), |
| 106 | ScheduledHyperParamSetter( |
| 107 | 'learning_rate', |
| 108 | [(0, BASE_LR), (30, BASE_LR * 1e-1), (60, BASE_LR * 1e-2), (80, BASE_LR * 1e-3)]), |
| 109 | DataParallelInferenceRunner( |
| 110 | dataset_val, infs, list(range(nr_tower))), |
| 111 | ] |
| 112 | |
| 113 | return TrainConfig( |
| 114 | model=Model(), |
| 115 | data=StagingInput(QueueInput(dataset_train)), |
| 116 | callbacks=callbacks, |
| 117 | steps_per_epoch=1281167 // total_batch, |
| 118 | max_epoch=100, |
| 119 | ) |
| 120 | |
| 121 | |
| 122 | if __name__ == '__main__': |
no test coverage detected