()
| 243 | assert input.setup_done() |
| 244 | |
| 245 | def get_grad_fn(): |
| 246 | ctx = get_current_tower_context() |
| 247 | inputs = input.get_input_tensors() |
| 248 | |
| 249 | def compute_grad_from_inputs(*inputs): |
| 250 | cost = get_cost_fn(*inputs) |
| 251 | assert isinstance(cost, tf.Tensor), \ |
| 252 | "Expect the given function to return a cost, but got {} instead".format(str(cost)) |
| 253 | assert cost.shape.ndims == 0, "Cost must be a scalar, but found {}!".format(cost) |
| 254 | |
| 255 | if not ctx.is_training: |
| 256 | return None # this is the tower function, could be called for inference |
| 257 | |
| 258 | if ctx.has_own_variables: |
| 259 | varlist = ctx.get_collection_in_tower(tfv1.GraphKeys.TRAINABLE_VARIABLES) |
| 260 | else: |
| 261 | varlist = tfv1.trainable_variables() |
| 262 | opt = get_opt_fn() |
| 263 | if is_tfv2() and isinstance(opt, tf.optimizers.Optimizer): |
| 264 | grads = opt.get_gradients(cost, varlist) |
| 265 | grads = list(zip(grads, varlist)) |
| 266 | else: |
| 267 | grads = opt.compute_gradients( |
| 268 | cost, var_list=varlist, |
| 269 | gate_gradients=self.GATE_GRADIENTS, |
| 270 | colocate_gradients_with_ops=self.COLOCATE_GRADIENTS_WITH_OPS, |
| 271 | aggregation_method=self.AGGREGATION_METHOD) |
| 272 | grads = FilterNoneGrad().process(grads) |
| 273 | return grads |
| 274 | |
| 275 | if not self.XLA_COMPILE: |
| 276 | return compute_grad_from_inputs(*inputs) |
| 277 | else: |
| 278 | try: |
| 279 | from tensorflow.contrib.compiler import xla # deprecated |
| 280 | except ImportError: |
| 281 | from tensorflow.python.compiler.xla import xla |
| 282 | |
| 283 | def xla_func(): |
| 284 | grads = compute_grad_from_inputs(*inputs) |
| 285 | # unpack, because the return value |
| 286 | # of xla function cannot have nested structure |
| 287 | grads = [x[0] for x in grads] |
| 288 | return grads |
| 289 | |
| 290 | grads_no_vars = xla.compile(xla_func) |
| 291 | if ctx.has_own_variables: |
| 292 | varlist = ctx.get_collection_in_tower(tfv1.GraphKeys.TRAINABLE_VARIABLES) |
| 293 | else: |
| 294 | varlist = tfv1.trainable_variables() |
| 295 | return list(zip(grads_no_vars, varlist)) |
| 296 | |
| 297 | return get_grad_fn |
nothing calls this directly
no test coverage detected