search

sksq96 / pytorch-summary

Model summary in PyTorch similar to `model.summary()` in Keras
MIT License
3.98k stars 412 forks source link

Save the model summary #164

Open arielsolomon opened 3 years ago

arielsolomon commented 3 years ago

Hi, How can I save the model summary output to variable? text file? print to file?

cainmagi commented 3 years ago

The current version could support your need. Please use the newest GitHub version instead of the version installed by PyPI, because torchsummary.summary_string does not exist in torchsummary 1.5.1.

Here is an example:

import torchsummary
module = ... # Get my torch.nn.Module
with open('my_module_log.log', 'w') as f:
    report, _ = torchsummary.summary_string(model, ...)
    f.write(report)