changing model to efficientnet

bayes/models.py

@@ -20,6 +20,7 @@ from sklearn.model_selection import train_test_split
 import torch
 from torchvision import models, transforms
 
+from efficientnet.tfkeras import EfficientNetB0
 from data.mnist.mnist_model import Net
 
 def get_xtrain(segs):
@@ -40,7 +41,7 @@ def get_xtrain(segs):
 def process_imagenet_get_model(data):
     """Gets wrapped imagenet model."""
     # Get the vgg16 model, used in the experiments
-    model = models.vgg16(pretrained=True)
+    model = EfficientNetB0(weights='imagenet')
     model.eval()
     # model.cuda()
 

requirements.txt

@@ -37,6 +37,7 @@ dice-ml
 distlib
 distributed
 drive-cli
+efficientnet
 entrypoints
 fasttext
 ffmpeg