add app

app.py

@@ -0,0 +1,72 @@
+import fasterai
+from fasterai.sparse.all import *
+from fasterai.prune.all import *
+import torch
+import gradio as gr
+import os
+from torch.ao.quantization import get_default_qconfig_mapping
+import torch.ao.quantization.quantize_fx as quantize_fx
+from torch.ao.quantization.quantize_fx import convert_fx, prepare_fx
+
+class Quant():
+    def __init__(self, backend="x86"):
+        self.qconfig = get_default_qconfig_mapping(backend)
+    
+    def quantize(self, model):
+        x = torch.randn(3, 224, 224)
+        model_prepared = prepare_fx(model.eval(), self.qconfig, x)
+        return convert_fx(model_prepared)
+
+
+def optimize_model(input_model, sparsity, context, criteria):
+
+    model = torch.load(input_model)
+    model = model.eval()
+    model = model.to('cpu')
+    sp = Sparsifier(model, 'filter', context, criteria=eval(criteria))
+    sp.sparsify_model(sparsity)
+    sp._clean_buffers()
+    pr = Pruner(model, context, criteria=eval(criteria))
+    pr.prune_model(sparsity)
+    qu = Quant()
+    qu_model = qu.quantize(model)
+    
+    comp_path = "./comp_model.pth"
+    
+    scripted = torch.jit.script(qu_model)
+    torch.jit.save(scripted, comp_path)
+    #torch.save(qu_model, comp_path)
+    
+    return comp_path
+
+def main_interface(model_file, sparsity, action):
+    if action == 'Speed':
+        return optimize_model(model_file, sparsity, 'local', "large_final")
+
+    if action == 'Size':
+        return optimize_model(model_file, sparsity, 'global', "large_final")
+
+    if action == 'Consumption':
+        return optimize_model(model_file, sparsity, 'local', "random")
+    else:
+        return "Action not supported"
+
+
+granularity = ['weight', 'filter']
+context = ['local', 'global']
+criteria = ['large_final', 'random']
+
+
+iface = gr.Interface(
+    fn=main_interface,
+    inputs= [            
+             gr.File(label="Upload your PyTorch model (.pth file)"),
+             gr.Slider(label="Compression Level", minimum=0, maximum=100),
+             gr.Radio(["Speed", "Size", "Consumption"], label="Select Action")  
+            ],
+    outputs=gr.File(label="Download Compressed Model"),
+    title="FasterAI",
+    description="Upload your neural network model (.pt file) and receive a compressed version.",
+)
+
+iface.launch()