Update app.py

app.py

@@ -4,10 +4,11 @@ import torch
 import torch.nn as nn
 import torchvision
 import torchvision.transforms as T
-
+from transformers import AutoFeatureExtractor, ResNetForImageClassification
 import timm
 
-model = timm.create_model("hf_hub:nateraw/resnet18-random", pretrained=True)
+feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/resnet-101")
+model = ResNetForImageClassification.from_pretrained("microsoft/resnet-101")
 model.eval()
 
 import os 
@@ -20,11 +21,12 @@ def print_bn():
             bn_data.extend(m.running_mean.data.numpy().tolist())
             bn_data.extend(m.running_var.data.numpy().tolist())
             bn_data.append(m.momentum)
+    print(len(bn_data))
     return bn_data
 
 def update_bn(image):
     cursor_im = 0
-    image = T.Resize((40,40))(image)
+    image = T.Resize((100,100))(image)
     image = image.reshape(-1)
     for m in model.modules():
         if(type(m) is nn.BatchNorm2d):
@@ -32,11 +34,11 @@ def update_bn(image):
                 M = m.running_mean.data.shape[0]
                 if(cursor_im+M < image.shape[0]):
                     m.running_mean.data = image[cursor_im:cursor_im+M]
-                    cursor_im += M # next
+                    cursor_im += M
                     print(cursor_im,':',cursor_im+M)
                 else:
                     m.running_mean.data[:image.shape[0]-cursor_im] = image[cursor_im:]
-                    break # finish 
+                    break 
     return
     
 
@@ -53,7 +55,7 @@ def greet(image):
         image = torch.permute(image, [0,3,1,2])
         update_bn(image)
         print(image.shape)
-        out = model(image)
+        out = model(pixel_values=image)
     return "Hello world!"