python - 在 Tensorflow 和范围中重新分配变量

Question

我试图理解为什么这些实现之一有效和另一个无效之间的区别。我正在尝试在 tensorflow 中表示一些几何图形。

首先，一个辅助文件，d_math.py

!/usr/bin/env python3

将 numpy 导入为 np将 tensorflow 导入为 tf

dtype = tf.float64

def skew_symmetric(vector):
    #Creates a tensorflow matrix which is a skew-symmetric version of the input vector    
    return tf.stack([(0., -vector[2], vector[1]), (vector[2], 0., -vector[0]), (-vector[1], vector[0], 0.)], axis=0)

这是实现 1:

#!/usr/bin/env python3
import numpy as np
import tensorflow as tf
import d_math as d
import math
import time


class Joint():
    def __init__(self, axis, pos): #TODO: right now only revolute:
        axis_ = tf.Variable(axis, dtype=d.dtype)
        axis_ /= tf.linalg.norm(axis)
        theta_ = tf.Variable(0.0, dtype=d.dtype) #Always at the 0 angle config
        self.theta_ = theta_
        self.R_ = tf.cos(theta_) * tf.eye(3, dtype=d.dtype) + d.skew_symmetric(axis_) + (1. - tf.cos(theta_)) * tf.einsum('i,j->ij', axis_, axis_)
        


joint = Joint(np.array([1.0, 1.0, 1.0]), 0.0)
init = tf.global_variables_initializer()    

with tf.Session() as session:
    session.run(init)    
    print(joint.R_)
    print(joint.R_.eval())
    joint.theta_ = joint.theta_.assign(math.pi/4.)
    session.run(joint.theta_)
    print(joint.R_.eval())
    

上面的版本更新了 theta，然后我得到了两个旋转矩阵的评估，一个用于 theta = 0，另一个用于 theta = pi/4。

然后，我尝试稍微重构我的代码，添加一个全局 session 变量，在单独的文件中创建，并在 API 中尽可能隐藏有关 tensorflow 的信息:

版本 2:

#!/usr/bin/env python3
import numpy as np
import tensorflow as tf
import d_math as d
import math
import time
import session as s


class Joint():
    def __init__(self, axis, pos): #TODO: right now only revolute:
        axis_ = tf.Variable(axis, dtype=d.dtype)
        axis_ = axis_ / tf.linalg.norm(axis)
        theta_ = tf.Variable(0.0, dtype=d.dtype) #Always at the 0 angle config
        self.theta_ = theta_
        self.R_ = tf.cos(theta_) * tf.eye(3, dtype=d.dtype) + d.skew_symmetric(axis_) + (1. - tf.cos(theta_)) * tf.einsum('i,j->ij', axis_, axis_)
        
    def set_theta(self, theta):
        self.theta_.assign(theta)
        s.session.run(self.theta_)
        

joint = Joint(np.array([1.0, 1.0, 1.0]), 0.0)
init = tf.global_variables_initializer()    

with s.session as session:
    session.run(init)  
    print(joint.R_)
    print(joint.R_.eval())
    #joint.theta_ = joint.theta_.assign(math.pi/4.)
    joint.set_theta(math.pi/4.)
    print(joint.R_.eval())

    

session.py 可以在这里看到:

#!/usr/bin/env python3
import tensorflow as tf

session = tf.Session()
    

这给出了两个评估的 R 矩阵，theta = 0。

有人可以向我解释为什么实现 2 不起作用吗？

Answer 1

tf.assign 返回更新变量的引用。根据文档:返回:一个张量，在分配完成后将保存“ref”的新值。

在第一个示例中，您实际上使用的是更新后的引用:

joint.theta_ = joint.theta_.assign(math.pi/4.)
session.run(joint.theta_)
print(joint.R_.eval())

在第二个示例中，您没有使用更新后的引用:

 def set_theta(self, theta):
    not_used = self.theta_.assign(theta)
    s.session.run(self.theta_)

我最好的猜测是，如果您使用更新的引用，它应该可以工作:

def set_theta(self, theta):
    self.theta_ = self.theta_.assign(theta)
    s.session.run(self.theta_)

此外，最好不要覆盖原始张量引用，因此我会为更新后的变量创建一个新属性:

def set_theta(self, theta):
    self.theta_updated_ = self.theta_.assign(theta)
    s.session.run(self.theta_updated_)

# ...
print(self.theta_updated_.eval())  # <<< This should give you updated value

重要:但是，运行 print(joint.R_.eval()) 可能无法为您提供更新的值，因为操作 self.R_ 不会强制依赖于更新后的引用 self.theta_updated_，您可能必须使用 tf.control_dependencies 来强制执行 self.R_ > 仅在更新完成后进行操作。例如:

with tf.control_dependencies([self.theta_updated_]):
    self.R_ = tf.cos(theta_) * # ...

最后注意:为变量赋值不会自动告诉其他操作它们需要等待该赋值完成。我艰难地发现了这一点。这里有一些snippets我写了跟踪使用 tf.assign 时变量的行为方式。我建议仔细阅读名为:优化已使用 tf.assign 更新的原始变量的代码片段。这些片段是独立的。