File size: 429 Bytes
477da44 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
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']
) |