PyTorch: 截断预训练ResNet50模型并替换Softmax层

import torch
import torch.nn as nn
import torchvision.models as models

# Load pre-trained ResNet50 model
resnet = models.resnet50(pretrained=True)

# Remove the last layer (softmax layer)
resnet = nn.Sequential(*list(resnet.children())[:-1])

# Define new softmax layer for your problem (30 classes)
num_classes = 30
new_softmax = nn.Linear(2048, num_classes)

# Combine ResNet50 and new softmax layer
model = nn.Sequential(resnet, new_softmax)

# Print the modified model architecture
print(model)

在这段代码中，我们首先导入所需的库。然后使用torchvision.models中的resnet50函数加载预训练的ResNet50模型。接下来，我们通过移除最后一层来截断模型。然后，我们定义一个新的softmax层，使用nn.Linear函数将ResNet50的输出大小（2048）映射到问题的类别数量（30）。最后，我们将截断的ResNet50和新的softmax层组合在一起，构建新的模型。

你可以根据需要修改num_classes的值来适应你的问题。