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
if __name__ == '__main__':
device = util.get_device()
# device = torch.device('cpu')
data_loaders, data_sizes = load_data()
print(data_loaders)
print(data_sizes)
res_loss = dict()
res_top1_acc = dict()
res_top5_acc = dict()
num_classes = 100
num_epochs = 50
for name in ['3e-4 -> 1e-4', '3e-4 -> 3e-5', '3e-4 -> 0']:
model = SqueezeNet(num_classes=num_classes)
model.eval()
# print(model)
model = model.to(device)
criterion = SmoothLabelCritierion(label_smoothing=0.1)
optimizer = optim.Adam(model.parameters(), lr=3e-4, weight_decay=3e-5)
if name == '3e-4 -> 1e-4':
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, num_epochs - 5, eta_min=1e-4)
elif name == '3e-4 -> 3e-5':
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, num_epochs - 5, eta_min=3e-5)
else:
lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, num_epochs - 5, eta_min=0)
warmup_scheduler = GradualWarmupScheduler(optimizer, multiplier=1, total_epoch=5, after_scheduler=lr_scheduler)
optimizer.zero_grad()
optimizer.step()
warmup_scheduler.step()