Failure to save weights and reuse them

[motavini]

Hello everyone!

I am currently working on solving the double-pendulum problem. After exploring the literature, I found that curriculum regularization (progressively increasing the complexity of the system during training) and sequence-to-sequence learning (dividing the time domain into small parts and training them one by one in sequence) appear to be effective techniques for tackling chaotic and high-frequency systems. To implement these strategies, I am attempting to save the weights obtained from a previous round of training and use them as initialization for a new round.

Unfortunately, I'm encountering difficulties in saving the weights after the training process. I have tried replicating some of the examples provided in the FAQ, but I have not been successful so far. I am running my code on Colab and I have included the relevant code below.

Additionally, I would appreciate guidance on how to initialize a new network with the saved weights and train it with different parameters in the ODE or on a different time domain.

import deepxde as dde
import numpy as np
import math
import tensorflow as tf

m1 = 1
m2 = 1
L1 = 1
L2 = 1
g = 9.81
theta1 = np.pi/180
theta1d = 0.0
theta2 = np.pi/360
theta2d = 0.0
def ode_system(x, y):

    y1, y2, y3, y4 = y[:, 0:1], y[:, 1:2], y[:, 2:3], y[:, 3:]
    dy1_x = dde.grad.jacobian(y, x, i=0)
    dy2_x = dde.grad.jacobian(y, x, i=1)
    dy3_x = dde.grad.jacobian(y, x, i=2)
    dy4_x = dde.grad.jacobian(y, x, i=3)
    eq1 = (-g*(2*m1+m2)*tf.sin(y1)-m2*g*tf.sin(y1-2*y2)-2*tf.sin(y1-y2)*m2*(L2*y4**2+L1*tf.cos(y1-y2)*y3**2))/(L1*(2*m1+m2-m2*tf.cos(2*y1-2*y2)))
    eq2 = (2*tf.sin(y1-y2)*(L1*(m1+m2)*y3**2+g*(m1+m2)*tf.cos(y1)+L2*m2*tf.cos(y1-y2)*y4**2))/(L2*(2*m1+m2-m2*tf.cos(2*y1-2*y2)))
    return [dy1_x - y3, dy2_x - y4, dy3_x - eq1, dy4_x - eq2]


def boundary(_, on_initial):
    return on_initial

geom = dde.geometry.TimeDomain(0, 8)
ic1 = dde.icbc.IC(geom, lambda x: theta1, boundary, component=0)
ic2 = dde.icbc.IC(geom, lambda x: theta2, boundary, component=1)
ic3 = dde.icbc.IC(geom, lambda x: theta1d, boundary, component=2)
ic4 = dde.icbc.IC(geom, lambda x: theta2d, boundary, component=3)
data = dde.data.PDE(geom, ode_system, [ic1, ic2, ic3, ic4], 1500, 2)

layer_size = [1] + [100] * 4 + [4]
activation = "tanh"
initializer = "Glorot uniform"
net1 = dde.nn.FNN(layer_size, activation, initializer)

def input_transform(t):
    return tf.concat(
        (
            t,
            tf.sin(t),
            tf.sin(2 * t),
            tf.sin(3 * t),
            tf.sin(4 * t),
            tf.sin(5 * t),
            tf.sin(6 * t),
            tf.sin(7 * t),
            tf.sin(8 * t),
            tf.sin(9 * t),
            tf.sin(10 * t),
            tf.sin(11 * t),
            tf.sin(12 * t),
            tf.sin(13 * t),
            tf.sin(14 * t),
            tf.sin(15 * t),
            tf.sin(16 * t),
            tf.sin(17 * t),
            tf.sin(18 * t),
            tf.sin(19 * t),
            tf.sin(20 * t),
            tf.sin(21 * t),
            tf.sin(22 * t),
            tf.sin(23 * t),
            tf.sin(24 * t),
            tf.sin(25 * t),
            tf.sin(34 * t),
        ),
        axis=1,
    )

net1.apply_feature_transform(input_transform)

model = dde.Model(data, net1)
model.compile("adam", lr=0.001)
losshistory, train_state = model.train(iterations=100)

model.compile("L-BFGS")
model.train_step.optimizer_kwargs = {'options': {'maxiter': 100}}

checker = dde.callbacks.ModelCheckpoint(
    "model/model.ckpt", save_better_only=True, period=50
)
model.train(epochs=100, callbacks=[checker])

dde.saveplot(losshistory, train_state, issave=True, isplot=True)