猫狗二分类添加一个拒绝类别，将猫狗预测概率小于阈值的预测样本归为第三类拒绝类别，训练和预测pytorh代码

以下基于PyTorch的猫狗二分类添加拒绝类别的训练和预测代码示例：

## 训练代码

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

# 定义数据预处理
transform = transforms.Compose(
[transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# 定义数据集
trainset = torchvision.datasets.ImageFolder(root='./train', transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=32, shuffle=True, num_workers=2)

# 定义模型
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 53 * 53, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 3) # 添加第三类拒绝类别

def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 53 * 53)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# 训练模型
for epoch in range(10): # 多次循环数据集
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
# 获取训练数据
inputs, labels = data

# 梯度清零
optimizer.zero_grad()

# 前向传播、计算损失、反向传播、优化
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()

# 打印统计信息
running_loss += loss.item()
if i % 200 == 199: # 每200个batch打印一次
print('[%d, %5d] loss: %.3f' %
(epoch + 1, i + 1, running_loss / 200))
running_loss = 0.0

print('Finished Training')

## 预测代码

# 定义测试数据集
testset = torchvision.datasets.ImageFolder(root='./test', transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2)

# 定义类别名称
classes = ('cat', 'dog', 'reject')

# 加载模型
net = Net()
net.load_state_dict(torch.load(PATH))

# 预测并输出结果
with torch.no_grad():
for data in testloader:
images, labels = data
outputs = net(images)

# 添加拒绝类别的判断
for i in range(len(outputs)):
probs = nn.functional.softmax(outputs[i], dim=0)
max_prob, predicted = torch.max(probs, 0)
if max_prob < threshold:
predicted = 2

print('Predicted: ', ' '.join('%s (%.2f)' % (classes[predicted[j]], probs[j]) for j in range(4)))