(cfg)
| 139 | |
| 140 | |
| 141 | def train(cfg): |
| 142 | startup_prog = fluid.Program() |
| 143 | train_prog = fluid.Program() |
| 144 | drop_last = True |
| 145 | dataset = build_dataset(cfg.DATASET.DATASET_NAME, |
| 146 | file_list=cfg.DATASET.TRAIN_FILE_LIST, |
| 147 | mode=ModelPhase.TRAIN, |
| 148 | shuffle=True, |
| 149 | data_dir=cfg.DATASET.DATA_DIR, |
| 150 | base_size= cfg.DATAAUG.BASE_SIZE, crop_size= cfg.DATAAUG.CROP_SIZE, rand_scale=True) |
| 151 | |
| 152 | def data_generator(): |
| 153 | if args.use_mpio: |
| 154 | data_gen = dataset.multiprocess_generator( |
| 155 | num_processes=cfg.DATALOADER.NUM_WORKERS, |
| 156 | max_queue_size=cfg.DATALOADER.BUF_SIZE) |
| 157 | else: |
| 158 | data_gen = dataset.generator() |
| 159 | |
| 160 | batch_data = [] |
| 161 | for b in data_gen: |
| 162 | batch_data.append(b) |
| 163 | if len(batch_data) == (cfg.TRAIN_BATCH_SIZE // cfg.NUM_TRAINERS): |
| 164 | for item in batch_data: |
| 165 | yield item[0], item[1], item[2] |
| 166 | batch_data = [] |
| 167 | # If use sync batch norm strategy, drop last batch if number of samples |
| 168 | # in batch_data is less then cfg.BATCH_SIZE to avoid NCCL hang issues |
| 169 | if not cfg.TRAIN.SYNC_BATCH_NORM: |
| 170 | for item in batch_data: |
| 171 | yield item[0], item[1], item[2] |
| 172 | |
| 173 | # Get device environment |
| 174 | gpu_id = int(os.environ.get('FLAGS_selected_gpus', 0)) |
| 175 | place = fluid.CUDAPlace(gpu_id) if args.use_gpu else fluid.CPUPlace() |
| 176 | places = fluid.cuda_places() if args.use_gpu else fluid.cpu_places() |
| 177 | |
| 178 | # Get number of GPU |
| 179 | dev_count = cfg.NUM_TRAINERS if cfg.NUM_TRAINERS > 1 else len(places) |
| 180 | print_info("#device count: {}".format(dev_count)) |
| 181 | cfg.TRAIN_BATCH_SIZE = dev_count * int(cfg.TRAIN_BATCH_SIZE_PER_GPU) |
| 182 | print_info("#train_batch_size: {}".format(cfg.TRAIN_BATCH_SIZE)) |
| 183 | print_info("#batch_size_per_dev: {}".format(cfg.TRAIN_BATCH_SIZE_PER_GPU)) |
| 184 | |
| 185 | py_reader, avg_loss, lr, pred, grts, masks = build_model( |
| 186 | train_prog, startup_prog, phase=ModelPhase.TRAIN) |
| 187 | py_reader.decorate_sample_generator( |
| 188 | data_generator, batch_size=cfg.TRAIN_BATCH_SIZE_PER_GPU, drop_last=drop_last) |
| 189 | |
| 190 | exe = fluid.Executor(place) |
| 191 | exe.run(startup_prog) |
| 192 | |
| 193 | exec_strategy = fluid.ExecutionStrategy() |
| 194 | # Clear temporary variables every 100 iteration |
| 195 | if args.use_gpu: |
| 196 | exec_strategy.num_threads = fluid.core.get_cuda_device_count() |
| 197 | exec_strategy.num_iteration_per_drop_scope = 100 |
| 198 | build_strategy = fluid.BuildStrategy() |
no test coverage detected