Build the GPUOptions proto.
(self)
| 896 | return config |
| 897 | |
| 898 | def _compute_gpu_options(self): |
| 899 | """Build the GPUOptions proto.""" |
| 900 | visible_device_list = [] |
| 901 | virtual_devices = [] |
| 902 | gpu_index = -1 |
| 903 | memory_growths = set() |
| 904 | for dev in self.list_physical_devices("GPU"): |
| 905 | gpu_index += 1 |
| 906 | |
| 907 | if dev not in self._visible_device_list: |
| 908 | continue |
| 909 | |
| 910 | growth = self._memory_growth_map[dev] |
| 911 | memory_growths.add(growth) |
| 912 | visible_device_list.append(str(gpu_index)) |
| 913 | |
| 914 | if self._virtual_device_map: |
| 915 | vdevs = self._virtual_device_map.get(dev, []) |
| 916 | device_limits = [] |
| 917 | for virt_dev in vdevs: |
| 918 | device_limits.append(virt_dev.memory_limit) |
| 919 | |
| 920 | virtual_devices.append( |
| 921 | config_pb2.GPUOptions.Experimental.VirtualDevices( |
| 922 | memory_limit_mb=device_limits)) |
| 923 | |
| 924 | # Only compute growth if virtual devices have not been configured and we |
| 925 | # have GPUs |
| 926 | if not virtual_devices and memory_growths: |
| 927 | if len(memory_growths) > 1: |
| 928 | raise ValueError("Memory growth cannot differ between GPU devices") |
| 929 | allow_growth = memory_growths.pop() |
| 930 | else: |
| 931 | allow_growth = None |
| 932 | |
| 933 | return config_pb2.GPUOptions( |
| 934 | allow_growth=allow_growth, |
| 935 | visible_device_list=",".join(visible_device_list), |
| 936 | experimental=config_pb2.GPUOptions.Experimental( |
| 937 | virtual_devices=virtual_devices)) |
| 938 | |
| 939 | @property |
| 940 | def function_call_options(self): |