风力发电机风扇语义分割数据集（包含python训练代码）

import numpy as np import torch from torch.utils.data import DataLoader from dataset import BrainMRIdataset from torchvision import transforms import matplotlib.pyplot as plt from Unet import UNet import torch.nn as nn import torch.optim as optim from tqdm import tqdm train_img_path = './data/train/image/' train_label_path = './data/train/label/' test_img_path = './data/test/image/' test_label_path = './data/test/label' train_transform = transforms.Compose([transforms.Resize((256,256)),transforms.ToTensor()]) test_transform = transforms.Compose([transforms.Resize((256,256)),transforms.ToTensor()]) # 加载数据集 trainSet = BrainMRIdataset(image_dir=train_img_path,mask_dir=train_label_path,transform=train_transform) testSet = BrainMRIdataset(image_dir=test_img_path,mask_dir=test_label_path,transform=test_transform) # 样本的个数 len_trainSet = len(trainSet) len_testSet = len(testSet) # 载入数据集 trainLoader = DataLoader(trainSet,batch_size=4,shuffle=True) testLoader = DataLoader(testSet,batch_size=4,shuffle=False) # 实例化模型 model =UNet(in_channels=3,num_classes=1) loss_fn = nn.BCEWithLogitsLoss() # 损失函数 optimizer = optim.Adam(model.parameters(),lr=0.0001) device = 'cuda' if torch.cuda.is_available() else 'cpu' model.to(device) # train def fit(epoch, model, train_loader, test_loader): correct = 0 running_loss = 0 epoch_IOU = [] # 训练模式 model.train() for x, y in tqdm(train_loader): x = x.to(device) y = y.to(device) y_pred = model(x) # forward loss = loss_fn(y_pred, y) # 损失函数 optimizer.zero_grad() # 梯度清零 loss.backward() # backward optimizer.step() # 梯度下降 with torch.no_grad(): y_pred = torch.sigmoid(y_pred) y_pred = (y_pred > 0.5).float() # 预测二值化 correct += (y_pred == y).sum().item() running_loss += loss.item() # 计算 iou intersection = y * y_pred # 计算交集 pred ∩ true temp = y_pred + y # pred + true union = temp - intersection # 计算并集：A ∪ B = A + B - A ∩ B batch_iou = intersection.sum() / (union.sum() + 1e-8) epoch_IOU.append(batch_iou.item()) epoch_loss = running_loss / len_trainSet epoch_acc = correct / (len_trainSet*256*256) test_correct = 0 test_running_loss = 0 epoch_test_IOU = [] model.eval() with torch.no_grad(): for x, y in tqdm(test_loader): x = x.to(device) y = y.to(device) y_pred = model(x) loss = loss_fn(y_pred, y) y_pred = torch.sigmoid(y_pred) y_pred = (y_pred > 0.5).float() # 预测二值化 test_correct += (y_pred == y).sum().item() test_running_loss += loss.item() # 计算 iou intersection = y * y_pred # 计算交集 pred ∩ true temp = y_pred + y # pred + true union = temp - intersection # 计算并集：A ∪ B = A + B - A ∩ B batch_test_iou = intersection.sum() / (union.sum() + 1e-8) epoch_test_IOU.append(batch_test_iou.item()) epoch_test_loss = test_running_loss / len_testSet epoch_test_acc = test_correct / (len_testSet*256*256) # save model static_dict = model.state_dict() torch.save(static_dict,'./log/pre_train_{}_train_iou_{}_test_iou_{}.pth'.format(epoch+1,round(np.mean(epoch_IOU),5),round(np.mean(epoch_test_IOU),5))) print('epoch:', epoch + 1) print('train loss:', round(epoch_loss, 5)) print('train accuracy:', round(epoch_acc, 5)) print('train IOU:',round(np.mean(epoch_IOU),5)) print('test loss:', round(epoch_test_loss, 5)) print('test accuracy:', round(epoch_test_acc, 5)) print('test IOU:', round(np.mean(epoch_test_IOU), 5)), if __name__ == '__main__': epochs = 50 for epoch in range(epochs): # 加载预训练权重 #model.load_state_dict(torch.load('./log/49_train_iou_0.86304_test_iou_0.81689.pth')) fit(epoch, model, train_loader=trainLoader,test_loader=testLoader) print('training over !!!')