tensorflow/tensorflow/python/framework/device.py at develop · NervanaSystems/tensorflow · GitHub

This repository was archived by the owner on Jan 3, 2023. It is now read-only.

History

320 lines (269 loc) · 9.89 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

# http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ==============================================================================

"""Class to represent a device."""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import copy

import threading

from tensorflow.python.util.tf_export import tf_export

@tf_export(v1=["DeviceSpec"])

class DeviceSpec(object):

"""Represents a (possibly partial) specification for a TensorFlow device.

`DeviceSpec`s are used throughout TensorFlow to describe where state is stored

and computations occur. Using `DeviceSpec` allows you to parse device spec

strings to verify their validity, merge them or compose them programmatically.

Example:

```python

# Place the operations on device "GPU:0" in the "ps" job.

device_spec = DeviceSpec(job="ps", device_type="GPU", device_index=0)

with tf.device(device_spec):

# Both my_var and squared_var will be placed on /job:ps/device:GPU:0.

my_var = tf.Variable(..., name="my_variable")

squared_var = tf.square(my_var)

```

If a `DeviceSpec` is partially specified, it will be merged with other

`DeviceSpec`s according to the scope in which it is defined. `DeviceSpec`

components defined in inner scopes take precedence over those defined in

outer scopes.

```python

with tf.device(DeviceSpec(job="train", )):

with tf.device(DeviceSpec(job="ps", device_type="GPU", device_index=0):

# Nodes created here will be assigned to /job:ps/device:GPU:0.

with tf.device(DeviceSpec(device_type="GPU", device_index=1):

# Nodes created here will be assigned to /job:train/device:GPU:1.

```

A `DeviceSpec` consists of 5 components -- each of

which is optionally specified:

* Job: The job name.

* Replica: The replica index.

* Task: The task index.

* Device type: The device type string (e.g. "CPU" or "GPU").

* Device index: The device index.

"""

def __init__(self, job=None, replica=None, task=None, device_type=None,

device_index=None):

"""Create a new `DeviceSpec` object.

Args:

job: string. Optional job name.

replica: int. Optional replica index.

task: int. Optional task index.

device_type: Optional device type string (e.g. "CPU" or "GPU")

device_index: int. Optional device index. If left

unspecified, device represents 'any' device_index.

"""

self.job = job

self.replica = replica

self.task = task

if device_type == "cpu" or device_type == "gpu":

# For backwards compatibility only, we support lowercase variants of

# cpu and gpu but turn them into uppercase here.

self.device_type = device_type.upper()

else:

self.device_type = device_type

self.device_index = device_index

self._hash = hash(self.to_string())

def _clear(self):

self._job = None

self._replica = None

self._task = None

self.device_type = None

self.device_index = None

@property

def job(self):

return self._job

@job.setter

def job(self, job):

if job is not None:

self._job = str(job)

else:

self._job = None

@property

def replica(self):

return self._replica

@replica.setter

def replica(self, replica):

if replica is not None:

self._replica = int(replica)

else:

self._replica = None

@property

def task(self):

return self._task

@task.setter

def task(self, task):

if task is not None:

self._task = int(task)

else:

self._task = None

def parse_from_string(self, spec):

"""Parse a `DeviceSpec` name into its components.

Args:

spec: a string of the form

/job:<name>/replica:<id>/task:<id>/device:CPU:<id>

/job:<name>/replica:<id>/task:<id>/device:GPU:<id>

as cpu and gpu are mutually exclusive.

All entries are optional.

Returns:

The `DeviceSpec`.

Raises:

ValueError: if the spec was not valid.

"""

self._clear()

splits = [x.split(":") for x in spec.split("/")]

for y in splits:

ly = len(y)

if y:

# NOTE(touts): we use the property getters here.

if ly == 2 and y[0] == "job":

self.job = y[1]

elif ly == 2 and y[0] == "replica":

self.replica = y[1]

elif ly == 2 and y[0] == "task":

self.task = y[1]

elif ((ly == 1 or ly == 2) and

((y[0].upper() == "GPU") or (y[0].upper() == "CPU"))):

if self.device_type is not None:

raise ValueError("Cannot specify multiple device types: %s" % spec)

self.device_type = y[0].upper()

if ly == 2 and y[1] != "*":

self.device_index = int(y[1])

elif ly == 3 and y[0] == "device":

if self.device_type is not None:

raise ValueError("Cannot specify multiple device types: %s" % spec)

self.device_type = y[1]

if y[2] != "*":

self.device_index = int(y[2])

elif ly and y[0] != "": # pylint: disable=g-explicit-bool-comparison

raise ValueError("Unknown attribute: '%s' in '%s'" % (y[0], spec))

return self

def merge_from(self, dev):

"""Merge the properties of "dev" into this `DeviceSpec`.

Args:

dev: a `DeviceSpec`.

"""

if dev.job is not None:

self.job = dev.job

if dev.replica is not None:

self.replica = dev.replica

if dev.task is not None:

self.task = dev.task

if dev.device_type is not None:

self.device_type = dev.device_type

if dev.device_index is not None:

self.device_index = dev.device_index

def to_string(self):

"""Return a string representation of this `DeviceSpec`.

Returns:

a string of the form

/job:<name>/replica:<id>/task:<id>/device:<device_type>:<id>.

"""

dev = ""

if self.job is not None:

dev += "/job:" + self.job

if self.replica is not None:

dev += "/replica:" + str(self.replica)

if self.task is not None:

dev += "/task:" + str(self.task)

if self.device_type is not None:

device_index_string = "*"

if self.device_index is not None:

device_index_string = str(self.device_index)

dev += "/device:%s:%s" % (self.device_type, device_index_string)

return dev

@staticmethod

def from_string(spec):

"""Construct a `DeviceSpec` from a string.

Args:

spec: a string of the form

/job:<name>/replica:<id>/task:<id>/device:CPU:<id>

/job:<name>/replica:<id>/task:<id>/device:GPU:<id>

as cpu and gpu are mutually exclusive.

All entries are optional.

Returns:

A DeviceSpec.

"""

return DeviceSpec().parse_from_string(spec)

def __eq__(self, other):

return self.to_string() == other.to_string()

def __hash__(self):

return self._hash

def check_valid(spec):

"""Check that a device spec is valid.

Args:

spec: a string.

Raises:

An exception if the spec is invalid.

"""

# Construct a DeviceSpec. It will assert a failure if spec is invalid.

DeviceSpec.from_string(spec)

def canonical_name(device):

"""Returns a canonical name for the given `DeviceSpec` or device name."""

if device is None:

return ""

if isinstance(device, DeviceSpec):

return device.to_string()

else:

device = DeviceSpec.from_string(device)

return device.to_string()

# Cache from DeviceSpec objects to their corresponding device functions.

# This cache is maintained for correctness, not performance: it makes it

# possible to compare the device function stacks belonging to different

# graphs in a meaningful way.

_cached_device_functions = {}

_cached_device_specs = {}

_cache_lock = threading.Lock()

def merge_device(spec):

"""Returns a device function that merges devices specifications.

This can be used to merge partial specifications of devices. The

innermost setting for a device field takes precedence. For example:

with tf.device(merge_device("/device:GPU:0"))

# Nodes created here have device "/device:GPU:0"

with tf.device(merge_device("/job:worker")):

# Nodes created here have device "/job:worker/device:GPU:0"

with tf.device(merge_device("/device:CPU:0")):

# Nodes created here have device "/job:worker/device:CPU:0"

with tf.device(merge_device("/job:ps")):

# Nodes created here have device "/job:ps/device:CPU:0"

Args:

spec: A `DeviceSpec` or a device spec string (partially) describing the

device that should be used for all nodes created in the scope of

the returned device function's with block.

Returns:

A device function with the above-described behavior.

Raises:

ValueError: if the spec was not valid.

"""

with _cache_lock:

if not isinstance(spec, DeviceSpec):

cached_device_spec = _cached_device_specs.get(spec, None)

if cached_device_spec is None:

device_spec = DeviceSpec.from_string(spec or "")

_cached_device_specs[spec] = device_spec

spec = device_spec

else:

spec = cached_device_spec

cached_function = _cached_device_functions.get(spec, None)

if cached_function is not None:

return cached_function

def _device_function(node_def):

current_device = DeviceSpec.from_string(node_def.device or "")

copy_spec = copy.copy(spec)

copy_spec.merge_from(current_device) # current_device takes precedence.

return copy_spec

_cached_device_functions[spec] = _device_function

return _device_function

tensorflow/tensorflow/python/framework/device.py at develop · NervanaSystems/tensorflow · GitHub

Uh oh!

Expand file tree

Breadcrumbs

device.py

Latest commit

History

Breadcrumbs

device.py

File metadata and controls