技术问询：如何复用已训练fastai模型及基于PyTorch加载其预训练模型

Hey there! Let's tackle your questions about reusing trained fastai models and loading them in a way similar to how you use pickle with scikit-learn. I'll break this down clearly for you.

Fastai makes reusing trained models straightforward, but you need to remember that the Learner object relies on matching the data structure you used during training. Here's a step-by-step walkthrough:

Saving the trained model #

Once your model is trained to your satisfaction, save its weights with:

learn.save('resnet34_test')  # This saves weights to the `models` folder in your PATH

Loading the model for reuse #

As you started doing, you first need to recreate the same data setup and Learner instance that matches the training environment:

from fastai.conv_learner import *

PATH = 'path/to/your/data'
sz = 224  # Same image size used during training
arch = resnet34

# Recreate the data object with identical parameters
data = ImageClassifierData.from_paths(PATH, tfms=tfms_from_model(arch, sz))
# Initialize the learner
learn = ConvLearner.pretrained(arch, data, precompute=False)
# Load your saved weights
learn.load('resnet34_test')

Making predictions #

To predict on new images, you'll use the validation transforms (val_tfms) to ensure consistency with how the model was trained:

# Get the training and validation transforms
trn_tfms, val_tfms = tfms_from_model(arch, sz)

# Process a single image
img = open_image('path/to/your/new_image.jpg')
# Apply validation transforms (avoid training transforms which include augmentation!)
img_transformed = val_tfms(img)
# Get predictions
preds = learn.predict(img_transformed)
# Map the prediction to the class name
predicted_class = data.classes[preds.argmax()]
print(f"Predicted class: {predicted_class}")

If you want something closer to scikit-learn's pickle workflow—where you can load the model without recreating the data setup manually—fastai has you covered with export() and load_learner(). This saves the entire Learner object (including data, model, and configuration):

Export the full learner #

After training, export the complete model package:

learn.export('my_trained_model.pkl')  # Saves to your PATH by default

Load the exported model #

Later, you can load it without rebuilding the data object first:

from fastai.conv_learner import load_learner

PATH = 'path/to/your/data'
# Load the entire learner
learn = load_learner(PATH, 'my_trained_model.pkl')

# Predict on new images directly
img = open_image('path/to/new_image.jpg')
preds = learn.predict(img)
predicted_class = learn.data.classes[preds.argmax()]
print(f"Predicted class: {predicted_class}")

If you want to work directly with PyTorch #

If you need to use the model purely within PyTorch (outside fastai), you can extract and save the PyTorch model weights:

# Save PyTorch model state dict
torch.save(learn.model.state_dict(), 'model_weights.pth')

# Later, load into a PyTorch model
from torchvision.models import resnet34
import torch.nn as nn

model = resnet34(pretrained=False)
# Adjust the final layer to match your number of classes
num_classes = len(data.classes)
model.fc = nn.Linear(model.fc.in_features, num_classes)
# Load the saved weights
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()  # Set to evaluation mode

A quick note: Make sure you're using the same fastai version when loading as you did when training—version mismatches can cause unexpected errors!

内容的提问来源于stack exchange，提问作者ni8mr