如何在PyTorch模型中保存embedding_stage层？附模型代码

Hey there! Let's break down how to save the embedding_stage (I noticed you had a typo as embeding_stage in your original code—fixed that in the examples below) based on your specific needs:

1. Save Only the Trainable Parameters for the Layer Behind embedding_stage #

Your embedding_stage is the result of F.relu(self.fc2(x)). Since ReLU has no trainable parameters, the actual learnable weights tied to this stage belong to the fc2 linear layer. You can save just these parameters like this:

model = model()
# Save fc2's parameters separately
torch.save(model.fc2.state_dict(), 'C:\\project\\count_speakers\\fc2_embedding_layer.h5')

To load them later:

# Initialize the fc2 layer (or the full model if needed)
fc2 = nn.Linear(128, 10)
fc2.load_state_dict(torch.load('C:\\project\\count_speakers\\fc2_embedding_layer.h5'))

2. Modify the Model to Return Both Predictions and embedding_stage #

If you want to keep using the full model but also access the embedding_stage when running inference, update the forward method to return both outputs:

import torch
import torch.nn as nn
import torch.nn.functional as F

class model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=12,out_channels=64,kernel_size=3,stride= 1,padding=1)
        # self.conv2 = nn.Conv2d(in_channels=64,out_channels=64,kernel_size=3,stride= 1,padding=1)
        self.fc1 = nn.Linear(24576, 128)
        self.bn = nn.BatchNorm1d(128)
        self.dropout1 = nn.Dropout2d(0.5)
        self.fc2 = nn.Linear(128, 10)
        self.fc3 = nn.Linear(10, 3)
    def forward(self, x):
        x = F.relu(self.conv1(x))
        # x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, (2,2))
        # print(x.shape)
        x = x.view(-1,24576)
        x = self.bn(F.relu(self.fc1(x)))
        x = self.dropout1(x)
        embedding_stage = F.relu(self.fc2(x))
        final_output = self.fc3(embedding_stage)
        return final_output, embedding_stage  # Return both results

Save the full model state_dict as you originally planned:

model = model()
torch.save(model.state_dict(), 'C:\\project\\count_speakers\\model_pytorch.h5')

When loading, you can now get both the prediction and embedding:

model = model()
model.load_state_dict(torch.load('C:\\project\\count_speakers\\model_pytorch.h5'))
model.eval()

# Example input (replace with your actual data)
sample_input = torch.randn(1, 12, 64, 64)  # Adjust shape to match your input

with torch.no_grad():
    pred, embedding = model(sample_input)
# `embedding` is your desired embedding_stage tensor

3. Save a Specific embedding_stage Tensor for a Given Input #

If you just need to save the actual embedding_stage tensor generated from a particular input (not the layer parameters), compute it directly and save:

model = model()
model.load_state_dict(torch.load('C:\\project\\count_speakers\\model_pytorch.h5'))
model.eval()

# Your input tensor (adjust shape as needed)
input_tensor = torch.randn(1, 12, 64, 64)

with torch.no_grad():
    # Replicate the forward pass up to the embedding_stage
    x = F.relu(model.conv1(input_tensor))
    x = F.max_pool2d(x, (2,2))
    x = x.view(-1,24576)
    x = model.bn(F.relu(model.fc1(x)))
    x = model.dropout1(x)
    embedding_stage = F.relu(model.fc2(x))

# Save the tensor
torch.save(embedding_stage, 'C:\\project\\count_speakers\\embedding_stage_tensor.pt')

To load the saved tensor later:

saved_embedding = torch.load('C:\\project\\count_speakers\\embedding_stage_tensor.pt')

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