Hmm… I am not sure what else I should put in place of
model = ModifiedSqueezeNetModel(num_classes = int(args.num_classes))


  
      gitlab.com
  
  
    promach/Pruning-CNN/blob/master/SqueezeNet-Pruning/finetune.py#L22


from prune import *
import argparse
import torch.nn.init as init
from operator import itemgetter
from heapq import nsmallest
import time




from flops import *


class ModifiedSqueezeNetModel(torch.nn.Module): # This class create an object 
	def __init__(self):
		super(ModifiedSqueezeNetModel, self).__init__()
		model = models.squeezenet1_1(pretrained=True)
		self.features = model.features
		for param in self.features.parameters():
			param.requires_grad = False


		final_conv = nn.Conv2d(512,2,kernel_size=1)
		self.classifier = nn.Sequential( 
			nn.Dropout(p=0.5),



  
  
    
    
  
  



[phung@archlinux SqueezeNet-Pruning]$ python predict.py --image “3_100.jpg” --model “model_prunned” --num_class “2”
Traceback (most recent call last):
File “predict.py”, line 23, in 
checkpoint = torch.load(path_to_model)
File “/usr/lib/python3.7/site-packages/torch/serialization.py”, line 358, in load
return _load(f, map_location, pickle_module)
File “/usr/lib/python3.7/site-packages/torch/serialization.py”, line 529, in _load
result = unpickler.load()
AttributeError: Can’t get attribute ‘ModifiedSqueezeNetModel’ on <module ‘main’ from ‘predict.py’>
[phung@archlinux SqueezeNet-Pruning]$

#Modified from https://github.com/amrit-das/Custom-Model-Training-PyTorch/blob/master/predict.py
import torch
import torch.nn as nn
#from torchvision.models import resnet18
from torchvision.transforms import transforms
import matplotlib.pyplot as plt
import numpy as np
from torch.autograd import Variable
import torch.functional as F
from PIL import Image
import os
import sys
import argparse
from prune import *

parser = argparse.ArgumentParser(description = 'To Predict from a trained model')
parser.add_argument('-i','--image', dest = 'image_name', required = True, help='Path to the image file')
parser.add_argument('-m','--model', dest = 'model_name', required = True, help='Path to the model')
parser.add_argument('-n','--num_class',dest = 'num_classes', required = True, help='Number of training classes')
args = parser.parse_args()

path_to_model = "./"+args.model_name
checkpoint = torch.load(path_to_model)

model = ModifiedSqueezeNetModel(num_classes = int(args.num_classes))
#model = resnet18(num_classes = int(args.num_classes))
model.load_state_dict(checkpoint)
model.eval()

def predict_image(image_path):
    print("prediciton in progress")
    image = Image.open(image_path)
    transformation = transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    image_tensor = transformation(image).float()
    image_tensor = image_tensor.unsqueeze_(0)

    if torch.cuda.is_available():
        image_tensor.cuda()

    input = Variable(image_tensor)
    output = model(input)

    index = output.data.numpy().argmax()
    return index

def class_mapping(index):
    mapping=open('class_mapping.txt','r')
    class_map={}
    for line in mapping:
        l=line.strip('\n').split('~')
        class_map[l[1]]=l[0]
    return class_map[str(index)]

if __name__ == "__main__":

    imagepath = "./test/Lemon/"+args.image_name
    prediction = predict_image(imagepath)
    name = class_mapping(prediction)
    print("Predicted Class: ",name)