使用Pytorch中的预训练模型进行语义分割，然后使用我们自己的数据集仅训练完全连接的图层

import torch
import torchvision

model = torchvision.models.segmentation.fcn_resnet50(pretrained=True)

num_classes = 2  # 根据你的数据集类别数进行设置
model.classifier[4] = torch.nn.Conv2d(512, num_classes, kernel_size=(1, 1))
model.aux_classifier[4] = torch.nn.Conv2d(256, num_classes, kernel_size=(1, 1))

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

num_epochs = 10  # 根据你的需求进行设置
for epoch in range(num_epochs):
    for images, labels in dataloader:  # dataloader是你准备的数据加载器
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs['out'], labels)
        loss.backward()
        optimizer.step()

model.eval()
with torch.no_grad():
    # 对测试集或新数据进行推理
    # ...