基于pytorch实现自定义图像数据集分类并绘制混淆矩阵、可视化散点图、损失曲线准确率曲线等代码实现
时间: 2023-12-18 09:03:25 浏览: 186
利用pytorch实现图像分类的一个完整的代码
好的,以下是一个基于PyTorch实现自定义图像数据集分类并绘制混淆矩阵、可视化散点图、损失曲线准确率曲线等的代码实现。请注意,这是一个基本的代码框架,您可能需要根据您的具体需求进行修改和完善。
```python
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from sklearn.metrics import confusion_matrix
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import numpy as np
# 定义超参数
num_epochs = 10
batch_size = 32
learning_rate = 0.001
# 定义数据预处理
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# 加载数据集
train_dataset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_dataset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# 定义模型
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
self.pool = nn.MaxPool2d(2, 2)
self.fc1 = nn.Linear(64 * 8 * 8, 512)
self.fc2 = nn.Linear(512, 10)
def forward(self, x):
x = self.pool(nn.functional.relu(self.conv1(x)))
x = self.pool(nn.functional.relu(self.conv2(x)))
x = x.view(-1, 64 * 8 * 8)
x = nn.functional.relu(self.fc1(x))
x = self.fc2(x)
return x
# 实例化模型和优化器
model = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9)
# 训练模型
train_loss = []
train_acc = []
for epoch in range(num_epochs):
running_loss = 0.0
running_corrects = 0
for i, (inputs, labels) in enumerate(train_loader):
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
_, preds = torch.max(outputs, 1)
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / len(train_dataset)
epoch_acc = running_corrects.double() / len(train_dataset)
train_loss.append(epoch_loss)
train_acc.append(epoch_acc)
print('Epoch [{}/{}], Loss: {:.4f}, Accuracy: {:.4f}'.format(epoch+1, num_epochs, epoch_loss, epoch_acc))
# 测试模型并绘制混淆矩阵
model.eval()
y_true = []
y_pred = []
with torch.no_grad():
for inputs, labels in test_loader:
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
y_true += labels.numpy().tolist()
y_pred += preds.numpy().tolist()
confusion = confusion_matrix(y_true, y_pred)
plt.imshow(confusion, cmap='binary')
plt.show()
# 可视化散点图
model.eval()
features = []
labels = []
with torch.no_grad():
for i, (inputs, labels_batch) in enumerate(test_loader):
outputs = model(inputs)
features_batch = outputs.numpy()
features.append(features_batch)
labels.append(labels_batch.numpy())
features = np.concatenate(features, axis=0)
labels = np.concatenate(labels, axis=0)
tsne = TSNE(n_components=2, random_state=0)
embeddings = tsne.fit_transform(features)
fig, ax = plt.subplots()
scatter = ax.scatter(embeddings[:,0], embeddings[:,1], c=labels)
legend = ax.legend(*scatter.legend_elements(), loc="upper right", title="Classes")
ax.add_artist(legend)
plt.show()
# 绘制损失曲线和准确率曲线
fig, ax = plt.subplots()
ax.plot(train_loss, label='train loss')
ax.plot(train_acc, label='train accuracy')
ax.set_xlabel('Epoch')
ax.legend()
plt.show()
```
希望这可以帮助您实现自定义图像数据集分类并绘制混淆矩阵、可视化散点图、损失曲线准确率曲线等。
阅读全文