| 75 | |
| 76 | |
| 77 | def Create(args): |
| 78 | gpus = list(range(args.num_gpus)) |
| 79 | log.info("Running on gpus: {}".format(gpus)) |
| 80 | |
| 81 | # Create CNNModeLhelper object |
| 82 | train_model = cnn.CNNModelHelper( |
| 83 | order="NCHW", |
| 84 | name="resnet50", |
| 85 | use_cudnn=True, |
| 86 | cudnn_exhaustive_search=False |
| 87 | ) |
| 88 | |
| 89 | # Model building functions |
| 90 | def create_resnet50_model_ops(model, loss_scale): |
| 91 | [softmax, loss] = resnet.create_resnet50( |
| 92 | model, |
| 93 | "data", |
| 94 | num_input_channels=3, |
| 95 | num_labels=1000, |
| 96 | label="label", |
| 97 | ) |
| 98 | model.Accuracy([softmax, "label"], "accuracy") |
| 99 | return [loss] |
| 100 | |
| 101 | # SGD |
| 102 | def add_parameter_update_ops(model): |
| 103 | model.AddWeightDecay(1e-4) |
| 104 | ITER = model.Iter("ITER") |
| 105 | stepsz = int(30) |
| 106 | LR = model.net.LearningRate( |
| 107 | [ITER], |
| 108 | "LR", |
| 109 | base_lr=0.1, |
| 110 | policy="step", |
| 111 | stepsize=stepsz, |
| 112 | gamma=0.1, |
| 113 | ) |
| 114 | AddMomentumParameterUpdate(model, LR) |
| 115 | |
| 116 | def add_image_input(model): |
| 117 | pass |
| 118 | |
| 119 | start_time = time.time() |
| 120 | |
| 121 | # Create parallelized model |
| 122 | data_parallel_model.Parallelize_GPU( |
| 123 | train_model, |
| 124 | input_builder_fun=add_image_input, |
| 125 | forward_pass_builder_fun=create_resnet50_model_ops, |
| 126 | param_update_builder_fun=add_parameter_update_ops, |
| 127 | devices=gpus, |
| 128 | ) |
| 129 | |
| 130 | ct = time.time() - start_time |
| 131 | train_model.net._CheckLookupTables() |
| 132 | |
| 133 | log.info("Model create for {} gpus took: {} secs".format(len(gpus), ct)) |
| 134 | |