unet脑肿瘤分割完整代码

import torch as t import torch.nn as nn from tqdm import tqdm #进度条 import segmentation_models_pytorch as smp from dataset import * device = t.device("cuda") if t.cuda.is_available() else t.device("cpu") train_data=BrainMRIdataset(train_img,train_label,train_transformer) test_data=BrainMRIdataset(test_img,test_label,test_transformer) dl_train=DataLoader(train_data,batch_size=4,shuffle=True) dl_test=DataLoader(test_data,batch_size=4,shuffle=True) model = smp.Unet( encoder_name="resnet34", # choose encoder, e.g. mobilenet_v2 or efficientnet-b7 encoder_weights="imagenet", # use `imagenet` pre-trained weights for encoder initialization in_channels=3, # model input channels (1 for gray-scale images, 3 for RGB, etc.) classes=2, # model output channels (number of classes in your dataset) ) img,label=next(iter(dl_train)) model=model.to('cuda') img=img.to('cuda') pred=model(img) label=label.to('cuda') loss_fn=nn.CrossEntropyLoss()#交叉熵损失函数 loss_fn(pred,label) optimizer=torch.optim.Adam(model.parameters(),lr=0.0001) def train_epoch(epoch, model, trainloader, testloader): correct = 0 total = 0 running_loss = 0 epoch_iou = [] #交并比 net=model.train() for x, y in tqdm(testloader): x, y = x.to('cuda'), y.to('cuda') y_pred = model(x) loss = loss_fn(y_pred, y) optimizer.zero_grad() loss.backward() optimizer.step() with torch.no_grad(): y_pred = torch.argmax(y_pred, dim=1) correct += (y_pred == y).sum().item() total += y.size(0) running_loss += loss.item() intersection = torch.logical_and(y, y_pred) union = torch.logical_or(y, y_pred) batch_iou = torch.sum(intersection) / torch.sum(union) epoch_iou.append(batch_iou.item()) epoch_loss = running_loss / len(trainloader.dataset) epoch_acc = correct / (total * 256 * 256) test_correct = 0 test_total = 0 test_running_loss = 0 epoch_test_iou = [] t.save(net.state_dict(), './Results2/weights/unet_weight/{}.pth'.format(epoch)) model.eval() with torch.no_grad(): for x, y in tqdm(testloader): x, y = x.to('cuda'), y.to('cuda') y_pred = model(x) loss = loss_fn(y_pred, y) y_pred = torch.argmax(y_pred, dim=1) test_correct += (y_pred == y).sum().item() test_total += y.size(0) test_running_loss += loss.item() intersection = torch.logical_and(y, y_pred)#预测值和真实值之间的交集 union = torch.logical_or(y, y_pred)#预测值和真实值之间的并集 batch_iou = torch.sum(intersection) / torch.sum(union) epoch_test_iou.append(batch_iou.item()) epoch_test_loss = test_running_loss / len(testloader.dataset) epoch_test_acc = test_correct / (test_total * 256 * 256)#预测正确的值除以总共的像素点 print('epoch: ', epoch, 'loss： ', round(epoch_loss, 3), 'accuracy:', round(epoch_acc, 3), 'IOU:', round(np.mean(epoch_iou), 3), 'test_loss： ', round(epoch_test_loss, 3), 'test_accuracy:', round(epoch_test_acc, 3), 'test_iou:', round(np.mean(epoch_test_iou), 3) ) return epoch_loss, epoch_acc, epoch_test_loss, epoch_test_acc if __name__ == "__main__": epochs=5 for epoch in range(epochs): train_epoch(epoch, model, dl_train, dl_test)