| import torch | |
| import torchvision.models as models | |
| from autosparsity.sparsity import sparsity_model | |
| if __name__ == "__main__": | |
| model = models.resnet50(pretrained=True).cuda() | |
| optimizer = None | |
| mode = 0 | |
| sparsity_model(model, optimizer, mode) | |
| model.eval() | |
| x = torch.randn((1,3,224,224)).cuda() | |
| torch.onnx.export( | |
| model, x, 'resnet50.onnx', input_names=['inputs'], output_names=['outputs'] | |
| ) |