ISAJ

1
@tf.function
2
def train_step(x, inverse_net, p_model, bg_maxs, bg_mins, soil_maxs, soil_mins,
3
               filling_rate_maxs, filling_rate_mins, input_pm):
4
    with tf.GradientTape() as tape:
5
        filling_rate, constant = inverse_net(x, training=True)
6
        n_filling = (filling_rate-filling_rate_mins) / (filling_rate_maxs-filling_rate_mins)
7
        n_soil = (input_pm-soil_mins) / (soil_maxs-soil_mins)
8
        p_input = tf.concat([n_soil, n_filling], axis=1)
9
        p_bg = p_model(p_input, training=False)
10
        p_bg = p_bg*(bg_maxs-bg_mins) + bg_mins
11
        p_bg = p_bg/constant
12
        loss = tf.reduce_mean(tf.square(p_bg[:, :2]-x))
13
    grads = tape.gradient(loss, inverse_net.trainable_variables)
14
    return loss, grads

1
optimizer.apply_gradients(zip(grads, inverse_net.trainable_variables))

1
loss, grads = train_step(input_bg, inverse_net, model, bg_maxs, bg_mins, soil_maxs, soil_mins,
2
                              filling_rate_maxs, filling_rate_mins, input_pm)
3
            optimizer.apply_gradients(zip(grads, inverse_net.trainable_variables))