Entry point into a DP thread :param gpu_idx: :param model: :param cluster_obj: :return:
(self, gpu_idx, model)
| 845 | logging.info(f'gpu available: {torch.cuda.is_available()}, used: {self.on_gpu}') |
| 846 | |
| 847 | def ddp_train(self, gpu_idx, model): |
| 848 | """ |
| 849 | Entry point into a DP thread |
| 850 | :param gpu_idx: |
| 851 | :param model: |
| 852 | :param cluster_obj: |
| 853 | :return: |
| 854 | """ |
| 855 | # otherwise default to node rank 0 |
| 856 | self.node_rank = 0 |
| 857 | |
| 858 | # show progressbar only on progress_rank 0 |
| 859 | self.show_progress_bar = self.show_progress_bar and self.node_rank == 0 and gpu_idx == 0 |
| 860 | |
| 861 | # determine which process we are and world size |
| 862 | if self.use_ddp: |
| 863 | self.proc_rank = self.node_rank * self.num_gpus + gpu_idx |
| 864 | self.world_size = self.num_gpus |
| 865 | |
| 866 | # let the exp know the rank to avoid overwriting logs |
| 867 | if self.logger is not None: |
| 868 | self.logger.rank = self.proc_rank |
| 869 | |
| 870 | # set up server using proc 0's ip address |
| 871 | # try to init for 20 times at max in case ports are taken |
| 872 | # where to store ip_table |
| 873 | model.trainer = self |
| 874 | model.init_ddp_connection(self.proc_rank, self.world_size) |
| 875 | |
| 876 | # CHOOSE OPTIMIZER |
| 877 | # allow for lr schedulers as well |
| 878 | model.model = model.build_model() |
| 879 | if not self.testing: |
| 880 | self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers()) |
| 881 | |
| 882 | # MODEL |
| 883 | # copy model to each gpu |
| 884 | if self.distributed_backend == 'ddp': |
| 885 | torch.cuda.set_device(gpu_idx) |
| 886 | model.cuda(gpu_idx) |
| 887 | |
| 888 | # set model properties before going into wrapper |
| 889 | self.copy_trainer_model_properties(model) |
| 890 | |
| 891 | # override root GPU |
| 892 | self.root_gpu = gpu_idx |
| 893 | |
| 894 | if self.distributed_backend == 'ddp': |
| 895 | device_ids = [gpu_idx] |
| 896 | else: |
| 897 | device_ids = None |
| 898 | |
| 899 | # allow user to configure ddp |
| 900 | model = model.configure_ddp(model, device_ids) |
| 901 | |
| 902 | # continue training routine |
| 903 | self.run_pretrain_routine(model) |
| 904 |
nothing calls this directly
no test coverage detected