[dattri.task] add partial parameter support for target/loss function and their gradient func

[TheaperDeng]

Description

1. Motivation and Context

1. This PR supports some extensive requirements for calculating the loss/target function and their grad/hessian based on only a portion of the parameter (e.g., some specific layers).
2. The support should be completely wrapped inside the task abstraction.

2. Summary of the change

1. Support the layer_name for task.get_grad_target_func, task.get_target_func, task.get_grad_loss_func, and task.get_loss_func (previously it will raise a notimplemented error).

TODO: this is the first PR to support partial parameter in dattri. Next PRs will support this feature in high-level attributor APIs.

task = AttributionTask(loss_func=f, model=model, checkpoints=model.state_dict())

# only use the last layer's parameter
grad_func_partial = task.get_grad_loss_func(
    layer_name=["fc3.weight", "fc3.bias"],
    index=0,  # this is a special additional parameter, because some other parameters we need to copy from the checkpoint.
)
params_partial, _ = task.get_param(
    layer_name=["fc3.weight", "fc3.bias"],
    index=0,
)
gradient_partial = grad_func_partial(params_partial, data_pair)

# only use the full parameter
grad_func_full = task.get_grad_loss_func()
params_full, _ = task.get_param()
gradient_full = grad_func_full(params_full, data_pair)

# the last layer's parameter's gradient should be the same as the last few dim full parameter's gradient.
torch.all_close(gradient_full[-last_layer_index:], gradient_partial)

3. What tests have been added/updated for the change?

• [x] Unit test: Typically, this should be included if you implemented a new function/fixed a bug.

[charles-pyj]

LGTM as well. I tried the gradients on both linear and convolution layers.

[charles-pyj]

But do we want to also add merging for the parameters for the same layer as well? For example, if one wants to conduct attribution on the first linear layer, it would be handy if we can combine the linear1.weight and linear1.bias together and obtain the concatenated gradients other than getting the separate gradients first and concatenate them later in attribution calculation?

[TheaperDeng]

But do we want to also add merging for the parameters for the same layer as well? For example, if one wants to conduct attribution on the first linear layer, it would be handy if we can combine the linear1.weight and linear1.bias together and obtain the concatenated gradients other than getting the separate gradients first and concatenate them later in attribution calculation?

I thought about this as well. Since we are using the named_parameters().keys() as the valid choices for layer names. I do afraid some times the layer's parameter is not named as xxx.weight and xxx.bias, then it may bring some error.