20 鱼币
请根据以下多标签分类ResNet网络模型和DenseNet网络模型分别修改提供的训练代码train.py,预测代码predict.py和测试代码evaluate.py,以分别适应两个网络模型
ResNet网络模型代码如下: import torch
import torch.nn as nn
import torch.nn.functional as F
# ResNet网络架构
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, in_planes, planes, stride=1):
super(BasicBlock, self).__init__()
self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.shortcut = nn.Sequential()
if stride != 1 or in_planes != self.expansion*planes:
self.shortcut = nn.Sequential(
nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
nn.BatchNorm2d(self.expansion*planes)
)
def forward(self, x):
out = F.relu(self.bn1(self.conv1(x)))
out = self.bn2(self.conv2(out))
out += self.shortcut(x)
out = F.relu(out)
return out
class ResNet(nn.Module):
def __init__(self, block, num_blocks, num_classes):
super(ResNet, self).__init__()
self.in_planes = 64
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
self.linear = nn.Linear(512*block.expansion, num_classes)
def _make_layer(self, block, planes, num_blocks, stride):
strides = [stride] + [1]*(num_blocks-1)
layers = []
for stride in strides:
layers.append(block(self.in_planes, planes, stride))
self.in_planes = planes * block.expansion
return nn.Sequential(*layers)
def forward(self, x):
out = F.relu(self.bn1(self.conv1(x)))
out = self.layer1(out)
out = self.layer2(out)
out = self.layer3(out)
out = self.layer4(out)
out = F.avg_pool2d(out, 4)
out = out.view(out.size(0), -1)
out = self.linear(out)
out = torch.sigmoid(out)
return out
def ResNet18(num_classes):
return ResNet(BasicBlock, [2,2,2,2], num_classes)
DenseNet网络模型代码如下: import torch
import torch.nn as nn
import torch.nn.functional as F
# DenseNet网络架构
class BasicBlock(nn.Module):
def __init__(self, inplanes, outplanes):
super(BasicBlock, self).__init__()
self.conv1 = nn.Conv2d(inplanes, outplanes, kernel_size=3, stride=1, padding=1, bias=False)
self.bn1 = nn.BatchNorm2d(outplanes)
self.relu1 = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(outplanes, outplanes, kernel_size=3, stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(outplanes)
self.relu2 = nn.ReLU(inplace=True)
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu1(out)
out = self.conv2(out)
out = self.bn2(out)
out += residual
out = self.relu2(out)
return out
class DenseBlock(nn.Module):
def __init__(self, inplanes, outplanes, num_layers):
super(DenseBlock, self).__init__()
self.blocks = nn.ModuleList([BasicBlock(inplanes + i * outplanes, outplanes) for i in range(num_layers)])
def forward(self, x):
features = [x]
for block in self.blocks:
y = block(torch.cat(features,dim=1))
features.append(y)
return torch.cat(features,dim=1)
class TransitionToFeatureMap(nn.Sequential):
def __init__(self, inplanes, outplanes):
super(TransitionToFeatureMap, self).__init__(
nn.Conv2d(inplanes, outplanes, kernel_size=1),
nn.BatchNorm2d(outplanes),
nn.AvgPool2d(kernel_size=2, stride=2))
class DenseNet(nn.Module):
def __init__(self, num_classes, growth_rate, layers_num):
super(DenseNet, self).__init__()
self.growth_rate = growth_rate
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu1 = nn.ReLU(inplace=True)
self.layer1 = nn.Sequential(
nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
DenseBlock(64, growth_rate, layers_num[0]))
self.inplanes = 64 + growth_rate * layers_num[0]
self.layer2 = nn.Sequential(
TransitionToFeatureMap(self.inplanes, 128),
DenseBlock(128, growth_rate, layers_num[1]))
self.inplanes += growth_rate * layers_num[1]
self.layer3 = nn.Sequential(
TransitionToFeatureMap(self.inplanes, 256),
DenseBlock(256, growth_rate, layers_num[2]))
self.inplanes += growth_rate * layers_num[2]
self.layer4 = nn.Sequential(
TransitionToFeatureMap(self.inplanes, 512),
DenseBlock(512, growth_rate, layers_num[3]))
self.inplanes += growth_rate * layers_num[3]
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.fc = nn.Linear(self.inplanes, num_classes)
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu1(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
x = torch.sigmoid(x)
return x
def DenseNet121(num_classes):
return DenseNet(num_classes=num_classes,
growth_rate=32,
layers_num=[6,12,24,16])
训练代码train.py: import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision import transforms
from model import ResNet18
import matplotlib.pyplot as plt
# 随机种子
torch.manual_seed(0)
# 设置参数和数据路径
lr = 0.0001
num_epochs = 20
batch_size = 8
data_path = '/home/a504/mjw/Code/data_set/plantvillage_demo1'
def main():
# 数据处理
transform_train = transforms.Compose([transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor()])
train_dataset = ImageFolder(root=data_path + '/train/', transform=transform_train)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
transform_val = transforms.Compose([transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor()])
val_dataset = ImageFolder(root=data_path + '/val/', transform=transform_val)
val_loader = DataLoader(dataset=val_dataset, batch_size=batch_size, shuffle=False)
# 初始化模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = ResNet18(num_classes=len(train_dataset.classes)).to(device)
# define loss function and optimizer
criterion = nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
# 训练模型
train_losses = []
val_losses = []
for epoch in range(num_epochs):
# 训练
train_loss = 0.0
model.train()
for images, labels in train_loader:
images, labels = images.to(device), labels.to(device)
labels = nn.functional.one_hot(labels, num_classes=len(train_dataset.classes)).float()
optimizer.zero_grad()
# forward pass
outputs = model(images)
loss = criterion(outputs, labels)
# backward pass and optimize
loss.backward()
optimizer.step()
train_loss += loss.item() * images.size(0)
train_loss = train_loss / len(train_loader.dataset)
train_losses.append(train_loss)
# validation
val_loss = 0.0
model.eval()
with torch.no_grad():
for images, labels in val_loader:
images, labels = images.to(device), labels.to(device)
labels = nn.functional.one_hot(labels, num_classes=len(val_dataset.classes)).float()
outputs = model(images)
loss = criterion(outputs, labels)
val_loss += loss.item() * images.size(0)
val_loss = val_loss / len(val_loader.dataset)
val_losses.append(val_loss)
print('Epoch [{}/{}], Train Loss: {:.4f}, Val Loss: {:.4f}'.
format(epoch + 1, num_epochs, train_loss, val_loss))
# 保存训练数据
torch.save(model.state_dict(), 'resnet18_multi_label_classification.pth')
# visualize the training process
plt.plot(range(num_epochs), train_losses, '-b', label='train')
plt.plot(range(num_epochs), val_losses, '-r', label='validation')
plt.legend(loc='lower right')
plt.xlabel('epoch')
plt.ylabel('loss')
plt.show()
if __name__ == '__main__':
main()
预测代码predict.py: import torch
import cv2
import numpy as np
from model import ResNet18
import matplotlib.pyplot as plt
# 加载模型
num_classes = 4
model = ResNet18(num_classes)
model.load_state_dict(torch.load('vgg16_multi_label_classification.pth'))
model.eval()
# 加载并预处理图像
image_path = 'OIP-C.jpeg'
image = cv2.imread(image_path)
image = cv2.resize(image, (224, 224))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_normalized = (image / 255.0 - 0.5) / 0.5
input_tensor = torch.from_numpy(image_normalized.transpose(2, 0, 1)).float().unsqueeze(0)
# 预测和可视化预测结果
with torch.no_grad():
output = model(input_tensor).squeeze().numpy()
class_names = ['Apple___Apple_scab', 'Apple___Black_rot', 'Apple___Cedar_apple_rust', 'Apple___healthy']
# acquire top2 and probabilities
top_indices = output.argsort()[-2:][::-1]
top_probabilities = output[top_indices]
top_class_names = [class_names[i] for i in top_indices]
# 在图像上标记预测结果及概率
fig, ax = plt.subplots()
ax.imshow(image)
for idx, (class_name, prob) in enumerate(zip(top_class_names, top_probabilities)):
plt.text(5, 15 + 20 * idx, f'{class_names[idx]}: {prob:.2f}', fontsize=12, color='red')
plt.show()
测试代码evaluate.py: import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision import transforms
from tqdm import tqdm
from model import ResNet18
# 常量和路径
num_classes = 4
model_path = './vgg16_multi_label_classification.pth'
test_path = '/home/a504/mjw/Code/data_set/plantvillage_demo1/val'
# 数据
transform_test = transforms.Compose([transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor()])
test_dataset = ImageFolder(test_path, transform=transform_test)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
# 模型
model = ResNet18(num_classes)
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
model.eval()
# define loss function
criterion = nn.BCELoss()
# 处理模型
hamming_loss = []
accuracy = []
precision = []
recall = []
f_beta = []
with torch.no_grad():
for idx, (images, labels) in tqdm(enumerate(test_loader), total=len(test_loader)):
outputs = model(images)
predicted = (outputs > 0.5).float() # threshold = 0.5
hamming_loss.append((predicted != labels.float()).sum().item() / (num_classes * len(labels)))
tp = ((predicted == 1) & (labels.float() == 1)).sum().item()
fp = ((predicted == 1) & (labels.float() == 0)).sum().item()
tn = ((predicted == 0) & (labels.float() == 0)).sum().item()
fn = ((predicted == 0) & (labels.float() == 1)).sum().item()
accuracy.append((tp + tn) / (num_classes * len(labels)))
precision.append(tp / (tp + fp + 1e-7))
recall.append(tp / (tp + fn + 1e-7))
f_beta.append(((1 + 0.5 ** 2) * precision[-1] * recall[-1]) / ((0.5 ** 2) * precision[-1] + recall[-1] + 1e-7))
print("Hamming Loss: {:.4f}".format(sum(hamming_loss) / len(hamming_loss)))
print("Accuracy: {:.4f}".format(sum(accuracy) / len(accuracy)))
print("Precision: {:.4f}".format(sum(precision) / len(precision)))
print("Recall: {:.4f}".format(sum(recall) / len(recall)))
print("F-beta Score: {:.4f}".format(sum(f_beta) / len(f_beta)))
我来回答