def squeezenet1_1(pretrained=False, **kwargs):
if pretrained:
model = SqueezeNet(version=1.1, **kwargs)
model.load_state_dict(
torch.load('./checkpoint/inception_v3_google-1a9a5a14.pth'))
return model
return SqueezeNet(version=1.1, **kwargs)
class ImgClassificationSqueezeNet:
def __init__(self, label_file_path, weights_file_path):
#self.model = squeezenet1_1(pretrained=True)
if type(weights_file_path) != type('str'):
weights_file_path = weights_file_path.decode("utf-8")
if type(label_file_path) != type('str'):
label_file_path = label_file_path.decode("utf-8")
is_existed = os.path.exists(weights_file_path)
if not is_existed:
print("Cannot find weights file...")
return -1
print(label_file_path)
self.model = SqueezeNet(version=1.1)
self.model.load_state_dict(torch.load(weights_file_path))
self.model.eval()
self.transformation = transforms.Compose([
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
is_existed = os.path.exists(label_file_path)
if not is_existed:
print("Error: no labels file found ...")
return -1
labels = json.load(open(label_file_path))
self.class_map = {int(key): value for (key, value) in labels.items()}
print('In python: initialize successfully ...')
def predict(self, image):
print("In python: predict function ...")
model = self.model
transformation = self.transformation
class_map = self.class_map
# Preprocess
image_tensor = transformation(image).float()
image_tensor = image_tensor.unsqueeze_(0)
if torch.cuda.is_available():
print("Using GPU ...")
image_tensor.cuda()
# Turn the input into a Variable
input = Variable(image_tensor)
# Predict the class of the image
output = model(input)
index = output.data.numpy().argmax()
prediction = class_map[index]
print("In python: predict done!")
return prediction
def run(self, img_data):
cv2_img = cv2.cvtColor(img_data[0], cv2.COLOR_BGR2RGB)
pil_img = Image.fromarray(cv2_img)
res = self.predict(pil_img)
return res
# 加载为c, 存储于CPU
c = torch.load('a.pth', map_location=lambda storage, loc: storage)
# 加载为d, 存储于GPU0上
d = torch.load('a.pth', map_location={'cuda:1': 'cuda:0'})
# In[40]:
torch.set_default_tensor_type('torch.FloatTensor')
from torchvision.models import SqueezeNet
model = SqueezeNet()
# module的state_dict是一个字典
model.state_dict().keys()
# Module对象的保存与加载
torch.save(model.state_dict(), 'squeezenet.pth')
model.load_state_dict(torch.load('squeezenet.pth'))
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
torch.save(optimizer.state_dict(), 'optimizer.pth')
optimizer.load_state_dict(torch.load('optimizer.pth'))
all_data = dict(optimizer=optimizer.state_dict(),
model=model.state_dict(),
info=u'模型和优化器的所有参数')
torch.save(all_data, 'all.pth')
all_data = torch.load('all.pth')
all_data.keys()