人工智能-遥感-语义分割-PyTorch实现高分遥感语义分割（地物分类）

import argparse import time import os import json from dataset import RSDataset import sync_transforms from torch.utils.data import DataLoader import torch.nn as nn import torch from models.deeplabv3_version_1.deeplabv3 import DeepLabV3 as model1 from models.deeplabv3_version_2.deeplabv3 import DeepLabV3 as model2 from libs import average_meter, metric from torch.autograd import Variable import numpy as np from tqdm import tqdm from prettytable import PrettyTable import torchvision from torchvision import transforms from palette import colorize_mask from PIL import Image from collections import OrderedDict from tensorboardX import SummaryWriter def parse_args(): parser = argparse.ArgumentParser(description="RemoteSensingSegmentation by PyTorch") # dataset parser.add_argument('--dataset-name', type=str, default='five') parser.add_argument('--train-data-root', type=str, default='../data_5_6_21/train') parser.add_argument('--val-data-root', type=str, default='../data_5_6_21/val') parser.add_argument('--train-batch-size', type=int, default=32, metavar='N', help='batch size for training (default:16)') parser.add_argument('--val-batch-size', type=int, default=32, metavar='N', help='batch size for testing (default:16)') # output_save_path parser.add_argument('--experiment-start-time', type=str, default=time.strftime('%m-%d-%H:%M:%S', time.localtime(time.time()))) parser.add_argument('--save-pseudo-data-path', type=str, default='/root/data/others/yaoganbisai/pseudo_data') # augmentation parser.add_argument('--base-size', type=int, default=512, help='base image size') parser.add_argument('--crop-size', type=int, default=512, help='crop image size') parser.add_argument('--flip-ratio', type=float, default=0.5) parser.add_argument('--resize-scale-range', type=str, default='0.5, 2.0') # model parser.add_argument('--model', type=str, default='deeplabv3_version_1', help='model name') parser.add_argument('--backbone', type=str, default='resnet50', help='backbone name') parser.add_argument('--pretrained', action='store_true', default=True) parser.add_argument('--n-blocks', type=str, default='3, 4, 23, 3', help='') parser.add_argument('--output-stride', type=int, default=16, help='') parser.add_argument('--multi-grids', type=str, default='1, 1, 1', help='') parser.add_argument('--deeplabv3-atrous-rates', type=str, default='6, 12, 18', help='') parser.add_argument('--deeplabv3-no-global-pooling', action='store_true', default=False) parser.add_argument('--deeplabv3-use-deformable-conv', action='store_true', default=False) parser.add_argument('--no-syncbn', action='store_true', default=False, help='using Synchronized Cross-GPU BatchNorm') # criterion parser.add_argument('--class-loss-weight', type=list, default= # [0.007814952234152803, 0.055862295151291756, 0.029094606950899726, 0.03104357983254851, 0.22757710412943985, 0.19666243636646102, 0.6088052968747066, 0.15683966777104494, 0.5288489922602664, 0.21668940382940433, 0.04310240828376457, 0.18284053575941367, 0.571096349549462, 0.32601488184885147, 0.45384359272537766, 1.0]) # [0.007956167959807792, 0.05664417300631733, 0.029857031694750392, 0.03198534634969046, 0.2309102255169529, # 0.19627322641039702, 0.6074939752850792, 0.16196525436190998, 0.5396602408824741, 0.22346488456565283, # 0.04453628275090391, 0.18672995330033487, 0.5990724459491834, 0.33183887346397484, 0.47737597643193597, 1.0] [0.008728536232175135, 0.05870821984204281, 0.030766985878693004, 0.03295408432939304, 0.2399409412190348, 0.20305583055639448, 0.6344888568739531, 0.16440413437125656, 0.5372260524694122, 0.22310945250778813, 0.04659596810284655, 0.19246378709444723, 0.6087430986295436, 0.34431415558778183, 0.4718853977371564, 1.0]) # loss parser.add_argument('--loss-names', type=str, default='cross_entropy') parser.add_argument('--classes-weight', type=str, default=None) parser.add_argument('--momentum', type=float, default=0.9, metavar='M', help='momentum (default:0.9)') parser.add_argument('--weight-decay', type=float, default=0.0001, metavar='M', help='weight-decay (default:1e-4)') # optimizer parser.add_argument('--optimizer-name', type=str, default='Adadelta') # learning_rate parser.add_argument('--base-lr', type=float, default=0.1, metavar='M', help='') # environment parser.add_argument('--use-cuda', action='store_true', default=True, help='using CUDA training') parser.add_argument('--num-GPUs', type=int, default=2, help='numbers of GPUs') parser.add_argument('--num_workers', type=int, default=4) # validation parser.add_argument('--eval', action='store_true', default=False, help='evaluation only') parser.add_argument('--no-val', action='store_true', default=False) parser.add_argument('--best-kappa', type=float, default=0) parser.add_argument('--total-epochs', type=int, default=12, metavar='N', help='number of epochs to train (default: 120)') parser.add_argument('--start-epoch', type=int, default=0, metavar='N', help='start epoch (default:0)') parser.add_argument('--resume-path', type=str, default=None) args = parser.parse_args() directory = "work_dirs/%s/%s/%s/%s/" % (args.dataset_name, args.model, args.backbone, args.experiment_start_time) args.directory = directory if not os.path.exists(directory): os.makedirs(directory) config_file = os.path.join(directory, 'config.json') with open(config_file, 'w') as file: json.dump(vars(args), file, indent=4) if args.use_cuda: print('Numbers of GPUs:', args.num_GPUs) else: print("Using CPU") return args class DeNormalize(object): def __init__(self, mean, std): self.mean = mean self.std = std def __call__(self, tensor): for t, m, s in zip(tensor, self.mean, self.std): t.mul_(s).add_(m) return tensor class Trainer(object): def __init__(self, args): self.args = args resize_scale_range = [float(scale) for scale in args.resize_scale_range.split(',')] sync_transform = sync_transforms.Compose([ sync_transforms.RandomScale(args.base_size, args.crop_size, resize_scale_range), sync_transforms.RandomFlip(args.flip_ratio) ]) self.resore_transform = transforms.Compose([ DeNormalize([.485, .456, .406], [.229, .224, .225]), transforms.ToPILImage() ]) self.visualize = transforms.Compose([transforms.ToTensor()]) class_name = args.dataset_name if class_name == 'fifteen': from class_names import fifteen_classes if class_name == 'five': from class_names import five_classes self.train_dataset = RSDataset(class_name, root=args.train_data_root, mode='train', sync_transforms=sync_transform) self.train_loader = DataLoader(dataset=self.train_dataset, batch_size=args.train_batch_size, num_workers=args.num_workers, shuffle=True, drop_last=True) print('class names {}.'.format(self.train_dataset.class_names)) print('Number samples {}.'.format(len(self.train_dataset))) if not args.no_val: val_data_set = RSDataset(class_name, root=args.val_data_root, mode='val', sync_transforms=None) self.val_loader = DataLoader(dataset=val_data_set, batch_size=args.val_batch_size, num_workers=args.num_workers, shuffle=False, drop_last=True) self.num_classes = len(self.train_dataset.class_names) print("类别数：", self.num_classes) self.class_loss_weight = torch.Te