Resnet预训练模型使用教程

ResNet 预训练模型使用教程
加载预训练的ResNet模型
为了开始使用ResNet预训练模型，在Python环境中需先导入必要的库，比如torchvision.models中的resnet模块。对于ResNet-50而言，可以通过指定参数pretrained=True来加载已经在ImageNet上预先训练好的权重。
import torchvision.models as models
model = models.resnet50(pretrained=True)
登录后复制
数据准备与预处理
当采用自定义的数据集时，通常会借助于torchvision.datasets.ImageFolder类来进行数据读取，并配合一系列转换操作以满足模型输入的要求[^1]。这些转换可能涉及调整图片大小、裁剪以及标准化等步骤：
from torchvision import transforms, datasets

data_transforms = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

image_dataset = datasets.ImageFolder(root='path_to_your_data', transform=data_transforms)

dataloaders = torch.utils.data.DataLoader(image_dataset, batch_size=4, shuffle=True)
登录后复制
修改分类头（全连接层）
如果目标任务不同于原始的ImageNet分类，则需要修改最后一层——即全连接层，使其适应新的类别数量。例如，如果有N个不同的标签，则应设置该层输出维度为N:
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, N) # 假设有N个类别
登录后复制
训练过程配置
完成上述准备工作之后，就可以进入实际的训练环节了。这里推荐使用Adam作为优化算法之一，并设定适当的学习率和其他超参数以便更好地收敛:
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(num_epochs):
running_loss = 0.0
for inputs, labels in dataloaders:
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()

running_loss += loss.item()
登录后复制