(device_ids)
| 44 | |
| 45 | |
| 46 | def _check_balance(device_ids): |
| 47 | imbalance_warn = """ |
| 48 | There is an imbalance between your GPUs. You may want to exclude GPU {} which |
| 49 | has less than 75% of the memory or cores of GPU {}. You can do so by setting |
| 50 | the device_ids argument to DataParallel, or by setting the CUDA_VISIBLE_DEVICES |
| 51 | environment variable.""" |
| 52 | |
| 53 | dev_props = [torch.cuda.get_device_properties(i) for i in device_ids] |
| 54 | |
| 55 | def warn_imbalance(get_prop): |
| 56 | values = [get_prop(props) for props in dev_props] |
| 57 | min_pos, min_val = min(enumerate(values), key=operator.itemgetter(1)) |
| 58 | max_pos, max_val = max(enumerate(values), key=operator.itemgetter(1)) |
| 59 | if min_val / max_val < 0.75: |
| 60 | warnings.warn(imbalance_warn.format(device_ids[min_pos], device_ids[max_pos])) |
| 61 | return True |
| 62 | return False |
| 63 | |
| 64 | if warn_imbalance(lambda props: props.total_memory): |
| 65 | return |
| 66 | if warn_imbalance(lambda props: props.multi_processor_count): |
| 67 | return |
| 68 | |
| 69 | |
| 70 |
no test coverage detected