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()