stanfordnlp / pyvene

Stanford NLP Python Library for Understanding and Improving PyTorch Models via Interventions
http://pyvene.ai
Apache License 2.0
559 stars 50 forks source link

[Minor] Adding support for any two-input lambda interventions; ResNet tutorials (no config) #105

Closed frankaging closed 5 months ago

frankaging commented 5 months ago

Description

Currently, we expose the APIs in the way that prioritizes intervention serializations, which means we can only use intervention primitives defined in pyvene which limits customizable intervention types. Although people can define interventions as nn.Modules, it takes more efforts comparing to a simple function, or lambda expression.

This change introduces how to interact with pyvene with a one-line lambda expression-based intervention which can also do trainable interventions with more customizable formats. At this point, pyvene basically acts like a torch.hook manager only losing its ability to share and save. Yet, it can easily support those simple interventions that are designed to be customizable.

Testing Done

       grad_fn=<AddBackward0>)
./juice/scr/wuzhengx/pyvene/pyvene/models/interventions.py:421: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  mask_sigmoid = torch.sigmoid(self.mask / torch.tensor(self.temperature))
.........Removing testing dir ./test_output_dir_prefix-2891ac
Removing testing dir ./test_output_dir_prefix-169855
Removing testing dir ./test_output_dir_prefix-6f0ab4
Removing testing dir ./test_output_dir_prefix-be15ba
Removing testing dir ./test_output_dir_prefix-50bb90
.............
----------------------------------------------------------------------
Ran 63 tests in 44.468s

OK

Checklist:

frankaging commented 5 months ago

maybe eval() can also take care of lambda sharing? since they are just one linear?