code: #!/usr/bin/env python3
  # -*- coding: utf-8 -*-
  # @author: Wesley
  # @time: 2020-12-11 10:47
  
  import os
  import time
  import torch
  from torch import nn
  from models.dinknet34 import DinkNet34
  from loss import dice_bce_loss
  from models.unet import UNet
  from dataset import MyDataset
  from torch.utils.data import DataLoader
  
  img_path = r'E:\PyCharmProject\datasets\5k\train_set\JPEGImages'
  mask_path = r'E:\PyCharmProject\datasets\5k\train_set\SegmentationClass'
  val_img_path = r'E:\PyCharmProject\datasets\5k\validate_set\JPEGImages'
  val_mask_path = r'E:\PyCharmProject\datasets\5k\validate_set\SegmentationClass'
  log = './dinknet.txt'
  
  
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  batch_size_per = 16
  batch_size = batch_size_per * torch.cuda.device_count()
  epoch_limit = 10
  net = DinkNet34().to(device)
  net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
  
  weight = r'E:\PyCharmProject\Road-Detection\weights\dinknet34.pt'
  # if os.path.exists(weight):
      # net.load_state_dict(torch.load(weight))
  
  train_dataset = MyDataset(img_path, mask_path)
  val_dataset = MyDataset(val_img_path, val_mask_path)
  train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
  val_dataloader = DataLoader(val_dataset, batch_size=batch_size)
  
  adam = torch.optim.Adam(net.parameters(), lr=2e-4)
  sgd = torch.optim.SGD(net.parameters(), lr=0.01, momentum=0.9)
  
  loss_fun = dice_bce_loss()
  
  if __name__ == '__main__':
  
      epoch = 1
      log = open(log, 'w', encoding='utf-8')
      log.write('epoch' + '\t' + 'loss' + '\t' + 'pa' + '\t' + 'iou' + '\t' + 'precision' + '\n')
      log.flush()
      while epoch < 300:
          s_time = time.time()
          print('epoch - {} - training'.format(epoch))
          net.train()
          TP = FP = TN = FN = 0
          pa = 0
          iou = 0
          stop = 0
          flag = 0
          train_loss = 0
          batch = len(train_dataloader)
          for i, (img, mask) in enumerate(train_dataloader):
              img = img.to(device)
              mask = mask.to(device)
              out = net(img)
              loss = loss_fun(mask, out)
  
              adam.zero_grad()
              loss.backward()
              adam.step()
  
              if i % 10 == 0:
                  print('{}: {}/{} - loss: {}'.format(epoch, i, batch, loss.item()))
                  # torch.save(net.state_dict(), weight)
                  # print('save success')
              train_loss += loss.item()
          epoch_loss = train_loss / len(train_dataloader)
  
          e_time = time.time()
          print('epoch - {} - epoch_loss: {}'.format(epoch, epoch_loss))
          print('total-time: ', e_time - s_time)
          print('epoch - {} - evaluating'.format(epoch))
  
          net.eval()
          for img, mask in val_dataloader:
              img = img.to(device)
              mask = mask.to(device)
              with torch.no_grad():
                  pred = net(img)
              pred[pred >= 0.5] = 1
              pred[pred < 0.5] = 0
  
              TP += ((pred == 1) & (mask == 1)).cpu().sum().item()
              TN += ((pred == 0) & (mask == 0)).cpu().sum().item()
              FN += ((pred == 0) & (mask == 1)).cpu().sum().item()
              FP += ((pred == 1) & (mask == 0)).cpu().sum().item()
  
          pa = (TP + TN) / (TP + TN + FP + FN)
          precision = TP / (TP + FN)
          iou = TP / (TP + FP + FN)
  
          print('pa: ', pa)
          print('iou: ', iou)
          print('precision', precision)
          log.write(
              str(epoch) + '\t' + str(epoch_loss) + '\t' + str(pa) + '\t' + str(iou) + '\t' + str(precision) + '\n')
          log.flush()
  
          if iou > stop:
              stop = iou
              torch.save(net.state_dict(), weight)
              print("save success，iou updated to: {}".format(iou))
              flag = 0
          else:
              flag += 1
              print("pa为{},没有提升，参数未更新,iou为{},第{}次未更新".format(iou, stop, flag))
              if flag >= epoch_limit:
                  print("early stop at epoch {}, finally iou: {}".format(epoch, stop))
                  break
          epoch += 1
      log.close()
  
apis: torch.utils.data.DataLoader torch.cuda.is_available torch.no_grad torch.cuda.device_count