用于图像分割的代码源码

import argparse import os from collections import OrderedDict from glob import glob import albumentations import pandas as pd import torch import torch.backends.cudnn as cudnn import torch.nn as nn import torch.optim as optim import yaml from albumentations import Flip from albumentations.augmentations import transforms from albumentations.core.composition import Compose, OneOf from sklearn.model_selection import train_test_split from torch.optim import lr_scheduler from tqdm import tqdm import archs import losses from dataset import Dataset from metrics import iou_score from utils import AverageMeter, str2bool ARCH_NAMES = archs.__all__ LOSS_NAMES = losses.__all__ LOSS_NAMES.append('BCEWithLogitsLoss') """ 指定参数： --dataset dsb2018_96 --arch NestedUNet """ def parse_args(): parser = argparse.ArgumentParser() parser.add_argument('--name', default=None, help='model name: (default: arch+timestamp)') parser.add_argument('--epochs', default=100, type=int, metavar='N', help='number of total epochs to run') parser.add_argument('-b', '--batch_size', default=8, type=int, metavar='N', help='mini-batch size (default: 16)') # model parser.add_argument('--arch', '-a', metavar='ARCH', default='NestedUNet', choices=ARCH_NAMES, help='model architecture: ' + ' | '.join(ARCH_NAMES) + ' (default: NestedUNet)') parser.add_argument('--deep_supervision', default=False, type=str2bool) parser.add_argument('--input_channels', default=3, type=int, help='input channels') parser.add_argument('--num_classes', default=1, type=int, help='number of classes') parser.add_argument('--input_w', default=96, type=int, help='image width') parser.add_argument('--input_h', default=96, type=int, help='image height') # loss parser.add_argument('--loss', default='BCEDiceLoss', choices=LOSS_NAMES, help='loss: ' + ' | '.join(LOSS_NAMES) + ' (default: BCEDiceLoss)') # dataset parser.add_argument('--dataset', default='dsb2018_96', help='dataset name') parser.add_argument('--img_ext', default='.png', help='image file extension') parser.add_argument('--mask_ext', default='.png', help='mask file extension') # optimizer parser.add_argument('--optimizer', default='SGD', choices=['Adam', 'SGD'], help='loss: ' + ' | '.join(['Adam', 'SGD']) + ' (default: Adam)') parser.add_argument('--lr', '--learning_rate', default=1e-3, type=float, metavar='LR', help='initial learning rate') parser.add_argument('--momentum', default=0.9, type=float, help='momentum') parser.add_argument('--weight_decay', default=1e-4, type=float, help='weight decay') parser.add_argument('--nesterov', default=False, type=str2bool, help='nesterov') # scheduler parser.add_argument('--scheduler', default='CosineAnnealingLR', choices=['CosineAnnealingLR', 'ReduceLROnPlateau', 'MultiStepLR', 'ConstantLR']) parser.add_argument('--min_lr', default=1e-5, type=float, help='minimum learning rate') parser.add_argument('--factor', default=0.1, type=float) parser.add_argument('--patience', default=2, type=int) parser.add_argument('--milestones', default='1,2', type=str) parser.add_argument('--gamma', default=2/3, type=float) parser.add_argument('--early_stopping', default=-1, type=int, metavar='N', help='early stopping (default: -1)') parser.add_argument('--num_workers', default=0, type=int) config = parser.parse_args() return config def train(config, train_loader, model, criterion, optimizer,device="mps"): avg_meters = {'loss': AverageMeter(), 'iou': AverageMeter()} model.train() pbar = tqdm(total=len(train_loader)) for input, target, _ in train_loader: input = input.to(device) target = target.to(device) # compute output if config['deep_supervision']: outputs = model(input) loss = 0 for output in outputs: loss += criterion(output, target) loss /= len(outputs) iou = iou_score(outputs[-1], target) else: output = model(input) loss = criterion(output, target) iou = iou_score(output, target) # compute gradient and do optimizing step optimizer.zero_grad() loss.backward() optimizer.step() avg_meters['loss'].update(loss.item(), input.size(0)) avg_meters['iou'].update(iou, input.size(0)) postfix = OrderedDict([ ('loss', avg_meters['loss'].avg), ('iou', avg_meters['iou'].avg), ]) pbar.set_postfix(postfix) pbar.update(1) pbar.close() return OrderedDict([('loss', avg_meters['loss'].avg), ('iou', avg_meters['iou'].avg)]) def validate(config, val_loader, model, criterion,device="mps"): avg_meters = {'loss': AverageMeter(), 'iou': AverageMeter()} # switch to evaluate mode model.eval() with torch.no_grad(): pbar = tqdm(total=len(val_loader)) for input, target, _ in val_loader: input = input.to(device) target = target.to(device) # compute output if config['deep_supervision']: outputs = model(input) loss = 0 for output in outputs: loss += criterion(output, target) loss /= len(outputs) iou = iou_score(outputs[-1], target) else: output = model(input) loss = criterion(output, target) iou = iou_score(output, target) avg_meters['loss'].update(loss.item(), input.size(0)) avg_meters['iou'].update(iou, input.size(0)) postfix = OrderedDict([ ('loss', avg_meters['loss'].avg), ('iou', avg_meters['iou'].avg), ]) pbar.set_postfix(postfix) pbar.update(1) pbar.close() return OrderedDict([('loss', avg_meters['loss'].avg), ('iou', avg_meters['iou'].avg)]) def main(): config = vars(parse_args()) if config['name'] is None: if config['deep_supervision']: config['name'] = '%s_%s_wDS' % (config['dataset'], config['arch']) else: config['name'] = '%s_%s_woDS' % (config['dataset'], config['arch']) os.makedirs('models/%s' % config['name'], exist_ok=True) print('-' * 20) for key in config: print('%s: %s' % (key, config[key])) print('-' * 20) with open('models/%s/config.yml' % config['name'], 'w') as f: yaml.dump(config, f) # define loss function (criterion) if config['loss'] == 'BCEWithLogitsLoss': criterion = nn.BCEWithLogitsLoss().cuda()#WithLogits 就是先将输出结果经过sigmoid再交叉熵 else: criterion = losses.__dict__[config['loss']]().cuda() cudnn.benchmark = True # create model print("=> creating model %s" % config['arch']) model = archs.__dict__[config['arch']](config['num_classes'], config['input_channels'], config['deep_supervision']) if torch.backends.mps.is_available():