neuralmagic / sparseml

Libraries for applying sparsification recipes to neural networks with a few lines of code, enabling faster and smaller models
Apache License 2.0
2.05k stars 144 forks source link

Move Session Management to Top Level #2261

Closed Satrat closed 4 months ago

Satrat commented 5 months ago

Previously to create, retrieve, or reset a session we needed to import import sparseml.core.session as session_manager. This file is now a top level import so instead the functions can be called by

Examples

Reset sessions between runs:

import sparseml

oneshot(.....)
sparseml.reset_session()
oneshot(....)

You can also use the create_session context manager to do this cleanup automatically. See tests/sparseml/transformers/finetune/test_finetune.py:test_oneshot_then_finetune() for a full example of this

import sparseml

with sparseml.create_session():
   oneshot(...)

with sparseml.create_session():
   oneshot(...)

Getting info from the current active session:

import sparseml

session = sparseml.active_session()
active_model = session.state.model
session_recipe = session.lifecycle.recipe_container.compiled_recipe 
mgoin commented 5 months ago

My only suggestion is to change the recommended usage to be explicitly using import sparseml to be clear what the "session" is connected to i.e.


import sparseml

oneshot(.....)
sparseml.reset_session()
oneshot(....)

with sparseml.create_session():
   oneshot(...)