def main(): n_classes = 1000 model = PeleeNet(n_classes=n_classes, pretrained=False) model.eval() input_data = misc.imread('../../../data/cat.jpg') # 按照imagenet的图像格式预处理 input_data = imagenet_utils.imagenet_preprocess(input_data) # x = Variable(torch.randn(1, 3, 224, 224)) x = Variable(torch.FloatTensor(torch.from_numpy(input_data))) y = Variable(torch.LongTensor(np.ones(1, dtype=np.int))) # print(x.shape) start = time.time() pred = model(x) end = time.time() print("PeleeNet forward time:", end - start) imagenet_utils.get_imagenet_label(pred)
m.weight.data.fill_(1) m.bias.data.zero_() elif isinstance(m, nn.Linear): n = m.weight.size(1) m.weight.data.normal_(0, 0.01) m.bias.data.zero_() if __name__ == '__main__': image_height, image_width, image_channel = (224, 224, 3) input = misc.imread('../../../data/cat.jpg') # 按照imagenet的图像格式预处理 input = imagenet_utils.imagenet_preprocess(input) n_classes = 1000 model = MobileNetV2(n_classes=n_classes) model.eval() # 训练模型为gpu模型 # model.load_state_dict(torch.load(os.path.expanduser('~/Data/mobilenetv2.pth.tar'), map_location=lambda storage, loc: storage)) # x = Variable(torch.randn(1, image_channel, image_height, image_width)) x = Variable(torch.FloatTensor(torch.from_numpy(input))) y = Variable(torch.LongTensor(np.ones(1, dtype=np.int))) # print(x.shape) start = time.time() pred = model(x) end = time.time() print("MobileNetV2 forward time:", end-start) imagenet_utils.get_imagenet_label(pred)