Nathan12 commited on
Commit
7ab4f00
·
1 Parent(s): 9d80ed6
Files changed (1) hide show
  1. app.py +72 -0
app.py ADDED
@@ -0,0 +1,72 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import fasterai
2
+ from fasterai.sparse.all import *
3
+ from fasterai.prune.all import *
4
+ import torch
5
+ import gradio as gr
6
+ import os
7
+ from torch.ao.quantization import get_default_qconfig_mapping
8
+ import torch.ao.quantization.quantize_fx as quantize_fx
9
+ from torch.ao.quantization.quantize_fx import convert_fx, prepare_fx
10
+
11
+ class Quant():
12
+ def __init__(self, backend="x86"):
13
+ self.qconfig = get_default_qconfig_mapping(backend)
14
+
15
+ def quantize(self, model):
16
+ x = torch.randn(3, 224, 224)
17
+ model_prepared = prepare_fx(model.eval(), self.qconfig, x)
18
+ return convert_fx(model_prepared)
19
+
20
+
21
+ def optimize_model(input_model, sparsity, context, criteria):
22
+
23
+ model = torch.load(input_model)
24
+ model = model.eval()
25
+ model = model.to('cpu')
26
+ sp = Sparsifier(model, 'filter', context, criteria=eval(criteria))
27
+ sp.sparsify_model(sparsity)
28
+ sp._clean_buffers()
29
+ pr = Pruner(model, context, criteria=eval(criteria))
30
+ pr.prune_model(sparsity)
31
+ qu = Quant()
32
+ qu_model = qu.quantize(model)
33
+
34
+ comp_path = "./comp_model.pth"
35
+
36
+ scripted = torch.jit.script(qu_model)
37
+ torch.jit.save(scripted, comp_path)
38
+ #torch.save(qu_model, comp_path)
39
+
40
+ return comp_path
41
+
42
+ def main_interface(model_file, sparsity, action):
43
+ if action == 'Speed':
44
+ return optimize_model(model_file, sparsity, 'local', "large_final")
45
+
46
+ if action == 'Size':
47
+ return optimize_model(model_file, sparsity, 'global', "large_final")
48
+
49
+ if action == 'Consumption':
50
+ return optimize_model(model_file, sparsity, 'local', "random")
51
+ else:
52
+ return "Action not supported"
53
+
54
+
55
+ granularity = ['weight', 'filter']
56
+ context = ['local', 'global']
57
+ criteria = ['large_final', 'random']
58
+
59
+
60
+ iface = gr.Interface(
61
+ fn=main_interface,
62
+ inputs= [
63
+ gr.File(label="Upload your PyTorch model (.pth file)"),
64
+ gr.Slider(label="Compression Level", minimum=0, maximum=100),
65
+ gr.Radio(["Speed", "Size", "Consumption"], label="Select Action")
66
+ ],
67
+ outputs=gr.File(label="Download Compressed Model"),
68
+ title="FasterAI",
69
+ description="Upload your neural network model (.pt file) and receive a compressed version.",
70
+ )
71
+
72
+ iface.launch()