tensorflow - 求解偏微分方程时在google colab上计算缓慢

Question

我正在使用 google colab 来求解均质热方程。我之前使用 scipy 使用稀疏矩阵制作了一个程序，该程序可以工作到 N = 10（超参数），但我需要像 N = 4...1000 一样运行它，因此它无法在我的电脑上运行。因此，我将代码转换为 tensorflow，在这里我无法像在 sympy 中那样使用稀疏矩阵，但即使 GPU/TPU 计算也比我的电脑慢而且慢。我在代码中面临并需要解决方案的问题

1) tf.contrib 被删除，因此我必须使用旧版本的 tensorflow 来执行 odeint 函数。它在 2.0 中在哪里？2）如果计算可以用稀疏矩阵计算，那可能会很好，因为矩阵是三对角矩阵。我知道 sparse_dense_mul() 函数，但它返回密集张量，它不会完成这项工作。“func”函数应用与时间无关的边界条件，然后需要 (nxn) 与 (nX1) 的矩阵乘法，从而得到 (nX1) 与多个矩阵。

在没有我创建课程的情况下，程序运行得更快。

它也给了这个

WARNING: Logging before flag parsing goes to stderr.
W0829 09:12:24.415445 139855355791232 lazy_loader.py:50] 
The TensorFlow contrib module will not be included in TensorFlow 2.0.
For more information, please see:
  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md
  * https://github.com/tensorflow/addons
  * https://github.com/tensorflow/io (for I/O related ops)
If you depend on functionality not listed there, please file an issue.

W0829 09:12:24.645356 139855355791232 deprecation.py:323] From /usr/local/lib/python3.6/dist-packages/tensorflow/contrib/integrate/python/ops/odes.py:233: div (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Deprecated in favor of operator or tf.math.divide.

当我在 range(2, 10) 中运行循环代码并且 tqdm 不显示并且单元格一直运行但它在 (2, 5) 中运行良好并且 tqdm bar 确实出现时。

#find a way to use sparse matrices
class Heat:

    def __init__(self, N):

        self.N = N
        self.H = 1/N

        self.A = ts.to_dense(ts.SparseTensor(indices=[[0, 0], [0, 1]] + \
                 [[i, i+j] for i in range(1, N) for j in [-1, 0, 1]] +[[N, N-1], [N, N]],
                  values=self.H*np.array([1/3, 1/6] + [1/6, 2/3, 1/6]*(N-1) + [1/6, 1/3], dtype=np.float32),
                  dense_shape=(N+1, N+1 )))

        self.D = ts.to_dense(ts.SparseTensor(indices=[[0, 0], [0, 1]] + [[i, i+j] \
                 for i in range(1, N) for j in [-1, 0, 1]] +[[N, N-1], [N, N]],
                 values=N*np.array([1-(1), -1 -(-1)] + [-1, 2, -1]*(N-1) + [-1-(-1), 1-(1)], dtype=np.float32),
                 dense_shape=(N+1, N+1)))

        self.domain = tf.linspace(0.0, 1.0, N+1)



        def f(k):

            if k == 0:

                return (1 + math.pi**2)*(math.pi*self.H - math.sin(math.pi*self.H))/(math.pi**2*self.H)

            elif k == N:

                return -(1 + math.pi**2)*(-math.pi*self.H + math.sin(math.pi*self.H))/(math.pi**2*self.H)

            else:

                return -2*(1 + math.pi**2)*(math.cos(math.pi*self.H) - 1)*math.sin(math.pi*self.H*k)/(math.pi**2*self.H)


        self.F = tf.constant([f(k) for k in range(N+1)], shape=(N+1,), dtype=tf.float32)  #caution! shape changed caution caution 1, N+1(problem) is different from N+1,

        self.exact = tm.scalar_mul(scalar=np.exp(1), x=tf.sin(math.pi*self.domain))



    def error(self):

        return np.linalg.norm(self.exact.numpy() -  self.approx, 2)


    def func (self, y, t):
        y = tf.Variable(y)
        y = y[0].assign(0.0)
        y = y[self.N].assign(0.0)
        if self.N**2> 100:
            y_dash = tl.matvec(tf.linalg.inv(self.A), tl.matvec(a=tm.negative(self.D), b=y, a_is_sparse=True) + tm.scalar_mul(scalar=math.exp(t), x=self.F))  #caution! shape changed F is (1, N+1) others too
        else:
            y_dash = tl.matvec(tf.linalg.inv(self.A), tl.matvec(a=tm.negative(self.D), b=y) + tm.scalar_mul(scalar=math.exp(t), x=self.F))  #caution! shape changed F is (1, N+1) others too

        y_dash = tf.Variable(y_dash) #!!y_dash performs Hadamard product like multiplication not matrix-like multiplication;returns 2-D

        y_dash = y_dash[0].assign(0.0)
        y_dash = y_dash[self.N].assign(0.0)

        return y_dash    

    def algo_1(self):

        self.approx = tf.contrib.integrate.odeint(
            func=self.func,
            y0=tf.sin(tm.scalar_mul(scalar=math.pi, x=self.domain)),
            t=tf.constant([0.0, 1.0]),
            rtol=1e-06,
            atol=1e-12,
            method='dopri5',
            options={"max_num_steps":10**10},
            full_output=False,
            name=None
           ).numpy()[1]

    def algo_2(self):

        self.approx = tf.contrib.integrate.odeint_fixed(
                      func=self.func,
                      y0=tf.sin(tm.scalar_mul(scalar=math.pi, x=self.domain)),
                      t=tf.constant([0.0, 1.0]),
                      dt=tf.constant([self.H**2], dtype=tf.float32),
                      method='rk4',
                      name=None
                    ).numpy()[1]


df = pd.DataFrame(columns=["NumBasis", "Errors"])
Ns = [2**r for r in range(2, 10)]
l =[]
for i in tqdm_notebook(Ns):
    heateqn = Heat(i)
    heateqn.algo_1()
    l.append([i, heateqn.error()])
    df.append({"NumBasis":i, "Errors":heateqn.error()}, ignore_index=True)
    tf.keras.backend.clear_session()

标签： tensorflowgpugoogle-colaboratoryfinite-element-analysis