could not import zoo

[mammadmaheri7]

Hi,

I encounter an error when I try to import ScheduledModifierManager:

versions:
sparsezoo-nightly: 1.5.0.20230531
sparseml-nightly: 1.5.4.20240501

code:
from sparseml.pytorch.optim import ScheduledModifierManager

log of error:
Traceback (most recent call last):
File "", line 1, in
File "/rds/general/user/mm6322/home/anaconda3/envs/CAP_new_new/lib/python3.9/site-packages/sparseml_nightly-1.5.4.20240501-py3.9.egg/sparseml/init.py", line 33, in
from .framework import (
File "/rds/general/user/mm6322/home/anaconda3/envs/CAP_new_new/lib/python3.9/site-packages/sparseml_nightly-1.5.4.20240501-py3.9.egg/sparseml/framework/init.py", line 22, in
from .info import *
File "/rds/general/user/mm6322/home/anaconda3/envs/CAP_new_new/lib/python3.9/site-packages/sparseml_nightly-1.5.4.20240501-py3.9.egg/sparseml/framework/info.py", line 55, in
from sparseml.sparsification.info import SparsificationInfo
File "/rds/general/user/mm6322/home/anaconda3/envs/CAP_new_new/lib/python3.9/site-packages/sparseml_nightly-1.5.4.20240501-py3.9.egg/sparseml/sparsification/init.py", line 22, in
from .analyzer import *
File "/rds/general/user/mm6322/home/anaconda3/envs/CAP_new_new/lib/python3.9/site-packages/sparseml_nightly-1.5.4.20240501-py3.9.egg/sparseml/sparsification/analyzer.py", line 26, in
from sparseml.optim import default_pruning_sparsities_loss
File "/rds/general/user/mm6322/home/anaconda3/envs/CAP_new_new/lib/python3.9/site-packages/sparseml_nightly-1.5.4.20240501-py3.9.egg/sparseml/optim/init.py", line 23, in
from .helpers import *
File "/rds/general/user/mm6322/home/anaconda3/envs/CAP_new_new/lib/python3.9/site-packages/sparseml_nightly-1.5.4.20240501-py3.9.egg/sparseml/optim/helpers.py", line 35, in
from sparsezoo import Zoo
ImportError: cannot import name 'Zoo' from 'sparsezoo' (/rds/general/user/mm6322/home/anaconda3/envs/CAP_new_new/lib/python3.9/site-packages/sparsezoo_nightly-1.5.0.20230531-py3.9.egg/sparsezoo/init.py)

[mgoin]

Hey @mammadmaheri7 , those versions are from our 1.5 nightly version which is quite old at this point, and the nightly versions don't match exactly. Could you try upgrading them to the latest version or forcing them to match the same dated version?

[mammadmaheri7]

As I trying to reproduce the result of recent published paper, they've implemented their codes based on older version of sparsezoo and updating sparsezoo will end up inconstancy. So upgrading sparsezoo is not possible.

[mammadmaheri7]

Also, I've tried latest version of sparsezoo (1.7.0) but still I encountered the same error:

trace:

Traceback (most recent call last):
File "one_shot_pruning.py", line 21, in
from sparseml.pytorch.optim import ScheduledModifierManager
File "/rds/general/user/mm6322/home/anaconda3/envs/CAP_final/lib/python3.8/site-packages/sparseml-1.7.0-py3.8.egg/sparseml/init.py", line 33, in
from .framework import (
File "/rds/general/user/mm6322/home/anaconda3/envs/CAP_final/lib/python3.8/site-packages/sparseml-1.7.0-py3.8.egg/sparseml/framework/init.py", line 22, in
from .info import *
File "/rds/general/user/mm6322/home/anaconda3/envs/CAP_final/lib/python3.8/site-packages/sparseml-1.7.0-py3.8.egg/sparseml/framework/info.py", line 55, in
from sparseml.sparsification.info import SparsificationInfo
File "/rds/general/user/mm6322/home/anaconda3/envs/CAP_final/lib/python3.8/site-packages/sparseml-1.7.0-py3.8.egg/sparseml/sparsification/init.py", line 22, in
from .analyzer import *
File "/rds/general/user/mm6322/home/anaconda3/envs/CAP_final/lib/python3.8/site-packages/sparseml-1.7.0-py3.8.egg/sparseml/sparsification/analyzer.py", line 26, in
from sparseml.optim import default_pruning_sparsities_loss
File "/rds/general/user/mm6322/home/anaconda3/envs/CAP_final/lib/python3.8/site-packages/sparseml-1.7.0-py3.8.egg/sparseml/optim/init.py", line 23, in
from .helpers import *
File "/rds/general/user/mm6322/home/anaconda3/envs/CAP_final/lib/python3.8/site-packages/sparseml-1.7.0-py3.8.egg/sparseml/optim/helpers.py", line 35, in
from sparsezoo import Zoo
ImportError: cannot import name 'Zoo' from 'sparsezoo' (/rds/general/user/mm6322/home/anaconda3/envs/CAP_final/lib/python3.8/site-packages/sparsezoo-1.7.0-py3.8.egg/sparsezoo/init.py)

[bfineran]

Hi @mammadmaheri7 can you maybe start from a clean environment? It looks like the import that it's failing on is quite old (even before 1.5) and is not up to date with what is in the 1.5 branch even. This also definitely should not be getting hit in 1.7 which leads me to believe there may be a local environment issue. Maybe it's anaconda? Our testing is done only through pip.

see the updated import on the 1.5 branch: https://github.com/neuralmagic/sparseml/blob/28894c3fcd87ea7cfca5101189d0ffb2c64a3e6f/src/sparseml/optim/helpers.py#L35