| 821 | return batch |
| 822 | |
| 823 | def set_distributed_mode(self, distributed_backend): |
| 824 | # skip for CPU |
| 825 | if self.num_gpus == 0: |
| 826 | return |
| 827 | |
| 828 | # single GPU case |
| 829 | # in single gpu case we allow ddp so we can train on multiple |
| 830 | # nodes, 1 gpu per node |
| 831 | elif self.num_gpus == 1: |
| 832 | self.single_gpu = True |
| 833 | self.use_dp = False |
| 834 | self.use_ddp = False |
| 835 | self.root_gpu = 0 |
| 836 | self.data_parallel_device_ids = [0] |
| 837 | else: |
| 838 | if distributed_backend is not None: |
| 839 | self.use_dp = distributed_backend == 'dp' |
| 840 | self.use_ddp = distributed_backend == 'ddp' |
| 841 | elif distributed_backend is None: |
| 842 | self.use_dp = True |
| 843 | self.use_ddp = False |
| 844 | |
| 845 | logging.info(f'gpu available: {torch.cuda.is_available()}, used: {self.on_gpu}') |
| 846 | |
| 847 | def ddp_train(self, gpu_idx, model): |
| 848 | """ |