@@ -123,7 +123,10 @@ def run_worker(self):
123123 gpus_env = os .getenv ("FLAGS_selected_gpus" )
124124 self .PSGPU = paddle .fluid .core .PSGPU ()
125125 gpuslot = [int (i ) for i in range (1 , self .model .sparse_inputs_slots )]
126+ gpu_mf_sizes = [self .model .sparse_feature_dim - 1 ] * (
127+ self .model .sparse_inputs_slots - 1 )
126128 self .PSGPU .set_slot_vector (gpuslot )
129+ self .PSGPU .set_slot_dim_vector (gpu_mf_sizes )
127130 self .PSGPU .init_gpu_ps ([int (s ) for s in gpus_env .split ("," )])
128131 opt_info = paddle .fluid .default_main_program ()._fleet_opt
129132 if use_auc is True :
@@ -139,7 +142,6 @@ def run_worker(self):
139142 if sync_mode == "heter" :
140143 self .heter_train_loop (epoch )
141144 elif sync_mode == "gpubox" :
142- self .reader ._set_use_ps_gpu (1 )
143145 self .dataset_train_loop (epoch )
144146 elif reader_type == "QueueDataset" :
145147 self .dataset_train_loop (epoch )
@@ -171,6 +173,7 @@ def run_worker(self):
171173 "Epoch: {}, using time {} second, ips {} {}/sec." .format (
172174 epoch , epoch_time , epoch_speed , self .count_method ))
173175 self .train_result_dict ["speed" ].append (epoch_speed )
176+ self .PSGPU .end_pass ()
174177
175178 model_dir = "{}/{}" .format (save_model_path , epoch )
176179 if fleet .is_first_worker (
@@ -181,7 +184,6 @@ def run_worker(self):
181184 self .inference_target_var )
182185 fleet .barrier_worker ()
183186 self .reader .release_memory ()
184- self .PSGPU .end_pass ()
185187 logger .info ("finish {} epoch training...." .format (epoch ))
186188 self .PSGPU .finalize ()
187189
0 commit comments