Блог

Есть ли более эффективный способ вычисления Якобиана (должен быть, он даже не выполняется для одной партии), я хочу вычислить потери, как указано в самоочевидной нейронной сети. Входные данные имеют форму (32, 365, 3), где 32-размер пакета. Потеря, которую я хочу свести к минимуму, — это уравнение 3 статьи.

Я считаю, что я не использую градиентную ленту оптимально.

 def compute_loss_theta(tape, parameter, concept, output, x):

    b = x.shape[0]
    in_dim = (x.shape[1], x.shape[2])

    feature_dim = in_dim[0]*in_dim[1]

    J = tape.batch_jacobian(concept, x)
    
    grad_fx = tape.gradient(output, x)
    grad_fx = tf.reshape(grad_fx,shape=(b, feature_dim))
    J = tf.reshape(J, shape=(b, feature_dim, feature_dim))

    parameter = tf.expand_dims(parameter, axis =1)

    loss_theta_matrix = grad_fx - tf.matmul(parameter, J)

    loss_theta = tf.norm(loss_theta_matrix)

    return loss_theta


for i in range(10):
    for x, y in train_dataset:

        with tf.GradientTape(persistent=True) as tape:
            tape.watch(x)
            
            parameter, concept, output = model(x)

            loss_theta = compute_loss_theta(tape, parameter, concept, output , x)

            loss_y = loss_object(y_true=y, y_pred=output)
                
            loss_value = loss_y   eps*loss_theta

        gradients = tape.gradient(loss_value, model.trainable_weights)
        optimizer.apply_gradients(zip(gradients, model.trainable_weights))