Register device specific flags. Args: tpu: Create flags to specify TPU operation. Returns: A list of flags for core.py to marks as key flags.
(tpu=True)
| 44 | |
| 45 | |
| 46 | def define_device(tpu=True): |
| 47 | """Register device specific flags. |
| 48 | |
| 49 | Args: |
| 50 | tpu: Create flags to specify TPU operation. |
| 51 | |
| 52 | Returns: |
| 53 | A list of flags for core.py to marks as key flags. |
| 54 | """ |
| 55 | |
| 56 | key_flags = [] |
| 57 | |
| 58 | if tpu: |
| 59 | flags.DEFINE_string( |
| 60 | name="tpu", |
| 61 | default=None, |
| 62 | help=help_wrap( |
| 63 | "The Cloud TPU to use for training. This should be either the name " |
| 64 | "used when creating the Cloud TPU, or a " |
| 65 | "grpc://ip.address.of.tpu:8470 url. Passing `local` will use the" |
| 66 | "CPU of the local instance instead. (Good for debugging.)")) |
| 67 | key_flags.append("tpu") |
| 68 | |
| 69 | flags.DEFINE_string( |
| 70 | name="tpu_zone", |
| 71 | default=None, |
| 72 | help=help_wrap( |
| 73 | "[Optional] GCE zone where the Cloud TPU is located in. If not " |
| 74 | "specified, we will attempt to automatically detect the GCE " |
| 75 | "project from metadata.")) |
| 76 | |
| 77 | flags.DEFINE_string( |
| 78 | name="tpu_gcp_project", |
| 79 | default=None, |
| 80 | help=help_wrap( |
| 81 | "[Optional] Project name for the Cloud TPU-enabled project. If not " |
| 82 | "specified, we will attempt to automatically detect the GCE " |
| 83 | "project from metadata.")) |
| 84 | |
| 85 | flags.DEFINE_integer( |
| 86 | name="num_tpu_shards", |
| 87 | default=8, |
| 88 | help=help_wrap("Number of shards (TPU chips).")) |
| 89 | |
| 90 | return key_flags |
nothing calls this directly
no test coverage detected