pytorch框架下加载图片做一次模型推理

if __name__ == "__main__":
    img1 = imread("/home/user/Downloads/112x112/re_bob.jpg")
    img2 = imread("/home/user/Downloads/112x112/re_jiwi.jpg")
    img1 = img1 / 255
    img2 = img2 / 255
    img1 = (img1 - 0.5) / 0.5
    img2 = (img2 - 0.5) / 0.5
    backbone = mx.mixnet_m(embedding_size=cfg.embedding_size, width_scale=cfg.scale, gdw_size=cfg.gdw_size).to('cuda:0')
    checkpoint = torch.load('./models/pretrained/mixfacenets-m_272928backbone.pth', map_location=torch.device('cuda:0'))
    backbone.load_state_dict(checkpoint, strict=False)

    model = torch.nn.DataParallel(backbone)
    model.eval()
    img1_embeddings = model(
        torch.tensor(np.expand_dims(np.moveaxis(img1, -1, 0), axis=0)).type(torch.float)).detach().cpu().numpy()
    img2_embeddings = model(
        torch.tensor(np.expand_dims(np.moveaxis(img2, -1, 0), axis=0)).type(torch.float)).detach().cpu().numpy()
    cos_sim = 1 - cdist(img1_embeddings, img2_embeddings, 'cosine')

    print('cos_sim=%f' % cos_sim)

你可能感兴趣的:(算法学习,pytorch,深度学习,python)