base_model_prep.py

import torch as th
from torchvision import models, transforms
from PIL import Image
import numpy as np
resnet18 = models.resnet18(pretrained=True)
# This stage below removes the last module (aka the softmax layer) to provide raw data for fine-tuning/transfer learning
modules=list(resnet18.children())[:-1]
model = th.nn.Sequential(*modules)

th.save(model, "/tmp/resnet18-headless.th")
# Lets check to see if the model was serialized properly
reloaded = th.load("/tmp/resnet18-headless.th")