takes a pytorch resnet 18 model, strips off the last module and checks to make sure serialization still works properly.

create_headless.py

import torch as th
from torchvision import models, transforms
from PIL import Image
import numpy as np
resnet18 = models.resnet18(pretrained=True)


modules=list(resnet18.children())[:-1]
model = th.nn.Sequential(*modules)

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
crop = transforms.CenterCrop(224)

image =  Image.open("/tmp/dog.jpg")
image_th = th.Tensor(np.asarray(crop(image)))
image_th = normalize(image_th)
image_th = image_th.reshape(1,3, 224, 224)

result = model.forward(image_th)
print(result)

th.save(model, "/tmp/resnet18-headless.th")
reloaded = th.load("/tmp/resnet18-headless.th")
newresult = reloaded.forward(image_th)
print(newresult)