变量简介

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 上查看源代码 下载笔记本

TensorFlow **变量** 是表示程序操作的共享持久状态的推荐方法。本指南介绍了如何在 TensorFlow 中创建、更新和管理 tf.Variable 实例。

变量通过 tf.Variable 类创建和跟踪。 tf.Variable 代表一个张量,其值可以通过对其运行操作来更改。特定的操作允许您读取和修改此张量的值。像 tf.keras 这样的高级库使用 tf.Variable 来存储模型参数。

设置

此笔记本讨论了变量放置。如果您想查看变量放置在哪个设备上,请取消注释此行。

import tensorflow as tf

# Uncomment to see where your variables get placed (see below)
# tf.debugging.set_log_device_placement(True)

创建变量

要创建变量,请提供初始值。 tf.Variable 将具有与初始化值相同的 dtype

my_tensor = tf.constant([[1.0, 2.0], [3.0, 4.0]])
my_variable = tf.Variable(my_tensor)

# Variables can be all kinds of types, just like tensors
bool_variable = tf.Variable([False, False, False, True])
complex_variable = tf.Variable([5 + 4j, 6 + 1j])

变量看起来和行为都像张量,实际上,它是一个由 tf.Tensor 支持的数据结构。与张量一样,它们具有 dtype 和形状,并且可以导出到 NumPy。

print("Shape: ", my_variable.shape)
print("DType: ", my_variable.dtype)
print("As NumPy: ", my_variable.numpy())

大多数张量操作按预期对变量起作用,尽管变量不能重新整形。

print("A variable:", my_variable)
print("\nViewed as a tensor:", tf.convert_to_tensor(my_variable))
print("\nIndex of highest value:", tf.math.argmax(my_variable))

# This creates a new tensor; it does not reshape the variable.
print("\nCopying and reshaping: ", tf.reshape(my_variable, [1,4]))

如上所述,变量由张量支持。您可以使用 tf.Variable.assign 重新分配张量。调用 assign 通常不会分配新的张量;而是重用现有张量的内存。

a = tf.Variable([2.0, 3.0])
# This will keep the same dtype, float32
a.assign([1, 2]) 
# Not allowed as it resizes the variable: 
try:
  a.assign([1.0, 2.0, 3.0])
except Exception as e:
  print(f"{type(e).__name__}: {e}")

如果您在操作中像张量一样使用变量,通常会对支持的张量进行操作。

从现有变量创建新变量会复制支持的张量。两个变量不会共享相同的内存。

a = tf.Variable([2.0, 3.0])
# Create b based on the value of a
b = tf.Variable(a)
a.assign([5, 6])

# a and b are different
print(a.numpy())
print(b.numpy())

# There are other versions of assign
print(a.assign_add([2,3]).numpy())  # [7. 9.]
print(a.assign_sub([7,9]).numpy())  # [0. 0.]

生命周期、命名和观察

在基于 Python 的 TensorFlow 中,tf.Variable 实例具有与其他 Python 对象相同的生命周期。当没有对变量的引用时,它会自动被释放。

变量也可以被命名,这可以帮助您跟踪和调试它们。您可以为两个变量赋予相同的名称。

# Create a and b; they will have the same name but will be backed by
# different tensors.
a = tf.Variable(my_tensor, name="Mark")
# A new variable with the same name, but different value
# Note that the scalar add is broadcast
b = tf.Variable(my_tensor + 1, name="Mark")

# These are elementwise-unequal, despite having the same name
print(a == b)

变量名称在保存和加载模型时会保留。默认情况下,模型中的变量会自动获取唯一的变量名称,因此您不需要自己分配它们,除非您需要。

虽然变量对于微分很重要,但有些变量不需要微分。您可以通过在创建时将 trainable 设置为 false 来关闭变量的梯度。一个不需要梯度的变量示例是训练步骤计数器。

step_counter = tf.Variable(1, trainable=False)

放置变量和张量

为了获得更好的性能,TensorFlow 会尝试将张量和变量放置在与其 dtype 兼容的最快的设备上。这意味着大多数变量都放置在 GPU 上(如果可用)。

但是,您可以覆盖它。在此代码段中,将浮点张量和变量放置在 CPU 上,即使 GPU 可用。通过打开设备放置日志记录(参见 设置),您可以看到变量放置的位置。

如果您在有和没有 GPU 的不同后端上运行此笔记本,您将看到不同的日志记录。请注意,必须在会话开始时打开设备放置日志记录。

with tf.device('CPU:0'):

  # Create some tensors
  a = tf.Variable([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
  c = tf.matmul(a, b)

print(c)

可以将变量或张量的地址设置为一个设备,并在另一个设备上进行计算。这将引入延迟,因为数据需要在设备之间复制。

但是,如果您有多个 GPU 工作器,但只想保留变量的一个副本,则可以这样做。

with tf.device('CPU:0'):
  a = tf.Variable([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
  b = tf.Variable([[1.0, 2.0, 3.0]])

with tf.device('GPU:0'):
  # Element-wise multiply
  k = a * b

print(k)

有关分布式训练的更多信息,请参阅 指南

下一步

要了解变量的典型使用方法,请参阅我们关于 自动微分 的指南。