pytorch 中草药识别代码

PyTorch是一种流行的深度学习框架，用于中药识别可以采用卷积神经网络（CNN）来处理图像数据。以下是一个简化的中药识别代码示例，假设我们已经有一个预处理好的中药图片数据集：

# 导入必要的库
import torch
import torchvision
from torchvision import datasets, transforms
import torch.nn as nn
import torch.optim as optim

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((64, 64)),  # 图像大小调整
    transforms.ToTensor(),        # 转换为张量
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 归一化
])

# 加载数据集（假设名为'drug_dataset'）
train_data = datasets.ImageFolder('path/to/drug_dataset/train', transform=transform)
test_data = datasets.ImageFolder('path/to/drug_dataset/test', transform=transform)

# 创建数据加载器
train_loader = torch.utils.data.DataLoader(train_data, batch_size=32, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=32, shuffle=False)

# 定义模型 - 这里是一个简单的卷积神经网络例子
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)  
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(32 * 64 * 64, 128)   
        self.fc2 = nn.Linear(128, num_classes)   # num_classes是中药类别数

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = x.view(-1, 32*64*64)     # Flatten the feature maps
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 初始化模型、损失函数和优化器
model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

# 测试模型
model.eval()
correct_count, total_count = 0, 0
with torch.no_grad():
    for images, labels in test_loader:
        predictions = model(images)
        _, predicted_labels = torch.max(predictions.data, 1)
        total_count += labels.size(0)
        correct_count += (predicted_labels == labels).sum().item()

accuracy = correct_count / total_count
print(f"Test accuracy: {accuracy}")