First, check there is a checkpoints/compressed_model_011000.pth in your workdir,
then create a empty folder fan_model, and create a new python file there fan_model/__init__.py
Finally, add a new python file just in the workdir
# load_compressed_FAN.py
import torch
from PIL import Image
import torchvision.transforms as transforms
# prepare input image
image = Image.open("./figure/eval_target_image.jpeg").convert('RGB')
totensor = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
x = totensor(image).unsqueeze(0)
x = x.cuda()
# load saved model with both structure and weights
# NOTICE: there must be ./fan_model/__init__.py which contains
# class fan_squeeze, Otherwise torch.load will fail
# but no need to import here, I don't know why
model = torch.load('./checkpoints/compressed_model_011000.pth')
model = model.cuda()
model.eval()
# run forward
with torch.no_grad():
print("x shape:", x.shape)
heatmap_68 = model(x)
print("heatmap shape:", heatmap_68.shape)
First, check there is a
checkpoints/compressed_model_011000.pth
in your workdir, then create a empty folderfan_model
, and create a new python file therefan_model/__init__.py
Finally, add a new python file just in the workdir
Now, check the dir structure:
run the
load_compressed_FAN.py
, it should work