AttributeError: module 'torch.onnx' has no attribute 'set_training' in pytorch1.6

[LovPe]

Hi Thanks for this job
I met a error when using pytorch1.6

  File "/home/adas/anaconda3/envs/py36_ch/lib/python3.6/site-packages/nni/compression/torch/speedup/compressor.py", line 59, in __init__
    self.torch_graph = build_module_graph(model, dummy_input)
  File "/home/adas/anaconda3/envs/py36_ch/lib/python3.6/site-packages/nni/_graph_utils.py", line 23, in build_module_graph
    return TorchModuleGraph(model, dummy_input)
  File "/home/adas/anaconda3/envs/py36_ch/lib/python3.6/site-packages/nni/_graph_utils.py", line 245, in __init__
    super().__init__(model, dummy_input, traced_model)
  File "/home/adas/anaconda3/envs/py36_ch/lib/python3.6/site-packages/nni/_graph_utils.py", line 65, in __init__
    self._trace(model, dummy_input)
  File "/home/adas/anaconda3/envs/py36_ch/lib/python3.6/site-packages/nni/_graph_utils.py", line 71, in _trace
    with torch.onnx.set_training(model, False):
AttributeError: module 'torch.onnx' has no attribute 'set_training'

I checked the doc in torch, it seems this has changed to torch.onnx.select_model_mode_for_export. Is there any plan to update this? Thanks

[AnshuTrivedi]

@LovPe can you share line of code which thrown error and complete error log with system info and libraries version?
It may help me to reproduce error and understand what went wrong.

[LovPe]

@AnshuTrivedi
Hi thanks for your relpay
You can use test code and the env info refer to #2879, this problem will happen.

[cvirus0626]

can you tell me the Method to deal with the error, Thank you very much.

[LovPe]

@cvirus0626 I did not handle this error since it is already in plan to solve in 2.0. Currently i suggest to use pytorch version < 1.6.

[bolgxh]

==========torchvision.version==========
0.7.0
==========torch.version==========
1.6.0
==========sys.version==========
3.7.9 (default, Aug 31 2020, 12:42:55)
[GCC 7.3.0]
I also have this porblem,when i test this code:

alexnet_model = torchvision.models.alexnet()
tw.draw_model(alexnet_model, [1, 3, 224, 224])

I think tensorwatch is a very good tools for deeplearning , I hope this problem can be resolved .

[kvartet]

This feature had been supported in PR #3076 and released in NNI v2.0. So I will close this issue, feel free to reopen it.

[LZhangGJ]

Thanks for this job. After version1.6 'set_training' may changed into 'select_model_mode_for_export'. I hope you can fix this problem because it is difficult to fix the sorce code in kaggle notebook. Thank you very much.