ndif-team / nnsight

The nnsight package enables interpreting and manipulating the internals of deep learned models.
https://nnsight.net/
MIT License
356 stars 34 forks source link

Early Exit of Sequential Graph Based Contexts #189

Closed AdamBelfki3 closed 1 month ago

AdamBelfki3 commented 1 month ago

New feature! You can now terminate early the execution of sequential graph-based contexts. The contexts concerned by this feature include the Session and Iterator contexts,Tracercontext does not support it. Call GraphBasedContext.exit() to skip the remaining execution.

Examples

Terminating a Session:

with model.session() as session:
    l = session.apply(list).save()

    l.append(0)
    l.append(1)
    session.exit()
    l.append(2)

print("Result: ", l)
"Result: [0, 1]"

Terminating an Iterator loop:

with model.session() as session:
    l = session.apply(list).save()

    with session.iter([0, 1, 2, 3]) as (item, iterator):
        with item == 2:
            iterator.exit()

        l.append(item)

print("Result: ", l)
"Result: [0, 1]"

Implementation

The Exit operation makes use of the EarlyStopProtocol, creating a node the moment it is called. The EarlyStopException is caught when the EarlyStopProtocol node is executed, with the Iterator breaking out of its loop.

Notes

An improvement should be put in place for the graph visualization feature to reflect sequential positioning of nodes in the graph.