如何训练自己的数据集之——太阳能光伏异常红外图像数据集

 import torch

    
 import torch.nn as nn
    
 import torch.optim as optim
    
 from torch.utils.data import Dataset, DataLoader
    
 from torchvision import transforms
    
 from PIL import Image
    
 import os
    
 import pandas as pd
    
 import numpy as np
    
  
    
 class SolarPanelDataset(Dataset):
    
     def __init__(self, root_dir, transform=None):
    
     self.root_dir = root_dir
    
     self.transform = transform
    
     self.labels = []
    
     self.images = []
    
  
    
     # 假设数据集有一个CSV文件来存储标签
    
     annotations_file = os.path.join(root_dir, 'annotations.csv')
    
     self.annotations = pd.read_csv(annotations_file)
    
     
    
     for index, row in self.annotations.iterrows():
    
         img_path = os.path.join(root_dir, row['filename'])
    
         image = Image.open(img_path)
    
         self.images.append(image)
    
         self.labels.append(row['label'])
    
  
    
     def __len__(self):
    
     return len(self.annotations)
    
  
    
     def __getitem__(self, idx):
    
     if torch.is_tensor(idx):
    
         idx = idx.tolist()
    
  
    
     img = self.images[idx]
    
     label = self.labels[idx]
    
  
    
     if self.transform:
    
         img = self.transform(img)
    
  
    
     return img, label
    
  
    
 # 数据变换
    
 transform = transforms.Compose([
    
     transforms.Resize((224, 224)),
    
     transforms.ToTensor(),
    
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    
 ])
    
  
    
 # 创建数据加载器
    
 root_dir = 'path/to/your/dataset'
    
 dataset = SolarPanelDataset(root_dir=root_dir, transform=transform)
    
 train_loader = DataLoader(dataset, batch_size=32, shuffle=True)
    
    
    
    
    AI写代码

 class SolarPanelClassifier(nn.Module):

    
     def __init__(self, num_classes=12):
    
     super(SolarPanelClassifier, self).__init__()
    
     self.features = nn.Sequential(
    
         nn.Conv2d(3, 64, kernel_size=3, padding=1),
    
         nn.ReLU(inplace=True),
    
         nn.MaxPool2d(kernel_size=2, stride=2),
    
  
    
         nn.Conv2d(64, 128, kernel_size=3, padding=1),
    
         nn.ReLU(inplace=True),
    
         nn.MaxPool2d(kernel_size=2, stride=2),
    
  
    
         nn.Conv2d(128, 256, kernel_size=3, padding=1),
    
         nn.ReLU(inplace=True),
    
         nn.MaxPool2d(kernel_size=2, stride=2),
    
  
    
         nn.Conv2d(256, 512, kernel_size=3, padding=1),
    
         nn.ReLU(inplace=True),
    
         nn.MaxPool2d(kernel_size=2, stride=2)
    
     )
    
  
    
     self.classifier = nn.Sequential(
    
         nn.Linear(512 * 7 * 7, 4096),
    
         nn.ReLU(True),
    
         nn.Dropout(),
    
         nn.Linear(4096, 4096),
    
         nn.ReLU(True),
    
         nn.Dropout(),
    
         nn.Linear(4096, num_classes)
    
     )
    
  
    
     def forward(self, x):
    
     x = self.features(x)
    
     x = x.view(x.size(0), -1)
    
     x = self.classifier(x)
    
     return x
    
    
    
    
    AI写代码

 # 初始化模型和优化器

    
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    
 model = SolarPanelClassifier().to(device)
    
 criterion = nn.CrossEntropyLoss()
    
 optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
    
  
    
 # 训练循环
    
 num_epochs = 20
    
 for epoch in range(num_epochs):
    
     running_loss = 0.0
    
     for i, (images, labels) in enumerate(train_loader, 0):
    
     images, labels = images.to(device), labels.to(device)
    
  
    
     # 零梯度
    
     optimizer.zero_grad()
    
  
    
     # 前向传播
    
     outputs = model(images)
    
     loss = criterion(outputs, labels)
    
  
    
     # 反向传播和优化
    
     loss.backward()
    
     optimizer.step()
    
  
    
     running_loss += loss.item()
    
     
    
     print(f'Epoch {epoch + 1}, Loss: {running_loss / len(train_loader)}')
    
  
    
 print('Finished Training')
    
    
    
    
    AI写代码

 correct = 0

    
 total = 0
    
 with torch.no_grad():
    
     for data in test_loader:
    
     images, labels = data[0].to(device), data[1].to(device)
    
     outputs = model(images)
    
     _, predicted = torch.max(outputs.data, 1)
    
     total += labels.size(0)
    
     correct += (predicted == labels).sum().item()
    
  
    
 print(f'Accuracy of the network on the test images: {100 * correct / total}%')
    
    
    
    
    AI写代码