[Bug]: - from_string constructor of "IHVPCalculator" doesn't work for "cgd" and "lissa"

Module

Common

Contact Details

No response

Current Behavior

When using, as expected, the from_string interface to IHVPCalculator, e.g. :

influence_calculator_cgd = FirstOrderInfluenceCalculator(influence_model, batched_ds, "cgd")

influence_calculator_lissa = FirstOrderInfluenceCalculator(influence_model, batched_ds, "lissa")

we get the following error :

TypeError: ConjugateGradientDescentIHVP.__init__() missing 1 required positional argument: 'train_dataset'

Note that this behavior isn't observed when using "exact" string.

Expected Behavior

Having train_dataset argument of ConjugateGradientDescentIHVP.__init__() and LissaIHVP.__init__() automatically set to batched_ds (that we already have to put in FirstOrderInfluenceCalculator

Version

v0.1.0

Environment

- OS:
- Python version:
- Tensorflow version:
- Packages used version:

Relevant log output

No response

To Reproduce

import numpy as np
np.random.seed(42)

import matplotlib.pyplot as plt
import seaborn as sns
sns.set_theme()

from sklearn.model_selection import train_test_split

import json

from tensorflow.keras.utils import to_categorical
import tensorflow as tf
tf.random.set_seed(42)

from deel.influenciae.common import InfluenceModel, ConjugateGradientDescentIHVP, ExactIHVP, LissaIHVP
from deel.influenciae.influence import FirstOrderInfluenceCalculator, SecondOrderInfluenceCalculator
from deel.influenciae.trac_in import TracIn
from deel.influenciae.utils import ORDER
from deel.influenciae.boundary_based import WeightsBoundaryCalculator
from deel.influenciae.benchmark.base_benchmark import ModelsSaver

def generate_data(N, center1 = np.array([-0.5,0]), center2 = np.array([0.5,0])):
  data = np.random.uniform(low=-1.,high=1.,size=(N,2))
  y = np.array([1*((x[0]-center1[0])**2 + (x[1]-center1[1])**2 <= 0.1 or (x[0]-center2[0])**2 + (x[1]-center2[1])**2 <= 0.1) for x in data])
  return data,y

def plot_data(x,y):
  for i,(pos_x,pos_y) in enumerate(x):
    if y[i]:
      plt.scatter(pos_x,pos_y,c='red',linewidths=0.5)
    else:
      plt.scatter(pos_x,pos_y,c='blue',linewidths=0.5)

x,y = generate_data(20)

X_train, X_test, y_train, y_test = train_test_split(x, y, test_size=0.33, random_state=42)

train_ds = tf.data.Dataset.from_tensor_slices((X_train, y_train)).map(lambda x, y: (tf.cast(x, tf.float32), tf.cast(y, tf.float32)))
test_ds = tf.data.Dataset.from_tensor_slices((X_test, y_test)).map(lambda x, y: (tf.cast(x, tf.float32), tf.cast(y, tf.float32)))

model = tf.keras.Sequential([
    tf.keras.layers.Dense(32, activation='relu'),
    tf.keras.layers.Dense(16, activation='relu'),
    tf.keras.layers.Dense(1, activation='sigmoid')
])

# optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)
# saver = ModelsSaver([3*i + 10 for i in range(14)],optimizer)

model.compile(loss='binary_crossentropy', optimizer="adam", metrics=['accuracy'])

model.fit(train_ds.batch(16), epochs=2, validation_data=test_ds.batch(16))

# Transform the Tensorflow model into an InfluenceModel (that implements some specific functionality for computing influences)

unreduced_loss = tf.keras.losses.BinaryCrossentropy(reduction=tf.keras.losses.Reduction.NONE)
influence_model = InfluenceModel(model, start_layer=-1, loss_function=unreduced_loss)

# Start by creating an object capable of efficiently computing inverse-hessian-vector products
# ihvp_calculator = ConjugateGradientDescentIHVP(influence_model, extractor_layer=-1,train_dataset=train_ds.batch(1))

# Create the influence calculator object
influence_calculator_exact = FirstOrderInfluenceCalculator(influence_model, train_ds.batch(1), "exact" )
influence_calculator_cgd = FirstOrderInfluenceCalculator(influence_model, train_ds.batch(1), "cgd")
influence_calculator_lissa = FirstOrderInfluenceCalculator(influence_model, train_ds.batch(1), "lissa")

deel-ai / influenciae