Faster RCNN 对肺结节 目标检测

from torchvision.transforms import transforms import torch import torchvision.models from dataset import LungDetection from torch.utils.data import DataLoader from utils import detection_collate, draw_bounding_box, miou import torch.optim as optim from tqdm import tqdm import numpy as np # 超参数设置 TRAIN_IMAGE_PATH = './data/train/image' # 训练集路径 TRAIN_LABEL_PATH = './data/train/label' # 训练集路径 TEST_IMAGE_PATH = './data/test/image' # 训练集路径 TEST_LABEL_PATH = './data/test/label' # 训练集路径 BATCH_SIZE, ROW, COLUMN = 4, 2, 2 # batch size + 可视化图像的行列数 LEARNING_RATE = 0.00001 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' EPOCHS = 50 # 预处理 transformer = transforms.Compose([transforms.ToTensor()]) # 加载训练集 trainSet = LungDetection(img_path=TRAIN_IMAGE_PATH,label_path=TRAIN_LABEL_PATH,transform=transformer) trainLoader = DataLoader(trainSet, BATCH_SIZE, shuffle=True, collate_fn=detection_collate) # 加载验证集 testSet = LungDetection(img_path=TEST_IMAGE_PATH,label_path=TEST_LABEL_PATH,transform=transformer) testLoader = DataLoader(testSet, batch_size=BATCH_SIZE, shuffle=False, collate_fn=detection_collate) # 建立模型 model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_fpn(weight=None, progress=True,num_classes=2) # model = torchvision.models.detection.retinanet_resnet50_fpn(pretrained=False, progress=True, num_classes=2) # model=torchvision.models.detection.ssd300_vgg16(pretrained=False, progress=True, num_classes=2) model.to(DEVICE) # 建立优化器 optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE) # 训练函数 def train_fn(model, optimizer, trainloader, testloader, device, epochs): for epoch in range(epochs): loss_epoch = [] # 一个epoch训练损失 iou_epoch = [] # 一个epoch训练 iou for images, labels in tqdm(trainloader): # 读取训练集 model.train() images = list(image.to(device) for image in images) # 将数据放到 cuda 上 labels = [{k: v.to(device) for k, v in t.items()} for t in labels] loss_dict = model(images, labels) # 前向传播 losses = sum(loss for loss in loss_dict.values()) # 损失求和 optimizer.zero_grad() # 梯度清零 losses.backward() # 反向传播 optimizer.step() # 梯度下降 losses = losses.data.cpu().numpy() loss_epoch.append(losses) # 保留损失 with torch.no_grad(): model.eval() # 测试模式 try: pred = model(images) batch_iou = miou(truth=labels, prediction=pred) # 计算一个batch的平均iou iou_epoch.append(batch_iou) except: continue test_loss_epoch = [] # 一个epoch 的 loss test_iou_epoch = [] # 一个epoch 的 iou with torch.no_grad(): # 验证的时候，不需要backward for images, labels in tqdm(testloader): model.train() # 训练模式 images = list(image.to(device) for image in images) labels = [{k: v.to(device) for k, v in t.items()} for t in labels] loss_dict = model(images, labels) losses = sum(loss for loss in loss_dict.values()) losses = losses.data.cpu().numpy() test_loss_epoch.append(losses) # 一个batch 的loss model.eval() # 测试模式 try: pred = model(images) test_batch_iou = miou(truth=labels, prediction=pred) # 计算一个batch的平均 iou test_iou_epoch.append(test_batch_iou) except: continue train_iou = np.array(iou_epoch).mean() # 将不同 batch 的 iou 求平均 train_loss = np.array(loss_epoch).mean() # 将不同 batch 的 loss 求均值 test_iou = np.array(test_iou_epoch).mean() test_loss = np.array(test_loss_epoch).mean() # 保留权重文件 if test_iou > 0.65: static_dict = model.state_dict() torch.save(static_dict, './log/pretrained_epoch_{},train_iou_{},test_iou_{}.pth'.format(epoch + 1, train_iou, test_iou)) print('epoch:', epoch + 1) print('epoch_train_loss:', train_loss) print('epoch_train_iou:', train_iou) print('epoch_test_loss:', test_loss) print('epoch_test_iou:', test_iou) if __name__ == '__main__': #model.load_state_dict(torch.load('./log/epoch_35,train_iou_0.9122412204742432,test_iou_0.6948733925819397.pth')) train_fn(model, optimizer, trainLoader, testLoader, DEVICE, EPOCHS) print(' training over !!!! ')