def test(args):
# output folder
outdir = 'outdir'
if not os.path.exists(outdir):
os.makedirs(outdir)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
if args.eval:
testset = get_segmentation_dataset(args.dataset,
split='val',
mode='testval',
transform=input_transform)
else:
testset = get_segmentation_dataset(args.dataset,
split='test',
mode='test',
transform=input_transform)
# dataloader
loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
test_data = data.DataLoader(testset,
batch_size=args.test_batch_size,
drop_last=False,
shuffle=False,
collate_fn=test_batchify_fn,
**loader_kwargs)
# model
if args.model_zoo is not None:
model = get_model(args.model_zoo, pretrained=True)
else:
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=BatchNorm2d,
base_size=args.base_size,
crop_size=args.crop_size)
# resuming checkpoint
if args.resume is None or not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
# strict=False, so that it is compatible with old pytorch saved models
model.load_state_dict(checkpoint['state_dict'])
utils.save_checkpoint({
'state_dict': checkpoint['state_dict'],
}, self.args, is_best, 'DANet101_reduce.pth.tar')
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
def test(args):
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# model
if args.model_zoo is not None:
model = get_model(args.model_zoo, pretrained=True)
else:
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
dilated=args.dilated,
lateral=args.lateral,
jpu=args.jpu,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=BatchNorm,
base_size=args.base_size,
crop_size=args.crop_size)
# resuming checkpoint
if args.resume is None or not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
# strict=False, so that it is compatible with old pytorch saved models
model.load_state_dict(checkpoint['state_dict'], strict=False)
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
print(model)
scales = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25] if args.dataset == 'citys' else \
[0.5, 0.75, 1.0, 1.25, 1.5, 1.75]
if not args.ms:
scales = [1.0]
num_classes = datasets[args.dataset.lower()].NUM_CLASS
evaluator = MultiEvalModule(model,
num_classes,
scales=scales,
flip=args.ms).cuda()
evaluator.eval()
img = input_transform(Image.open(
args.input_path).convert('RGB')).unsqueeze(0)
with torch.no_grad():
output = evaluator.parallel_forward(img)[0]
predict = torch.max(output, 1)[1].cpu().numpy()
mask = utils.get_mask_pallete(predict, args.dataset)
mask.save(args.save_path)
def __init__(self, args):
self.args = args
args.log_name = str(args.checkname)
root_dir = getattr(args, "data_root", '../datasets')
wo_head = getattr(args, "resume_wo_head", False)
self.logger = utils.create_logger(args.log_root, args.log_name)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size,
'logger': self.logger,
'scale': args.scale
}
trainset = get_segmentation_dataset(args.dataset,
split='train',
mode='train',
root=root_dir,
**data_kwargs)
testset = get_segmentation_dataset(args.dataset,
split='val',
mode='val',
root=root_dir,
**data_kwargs)
# dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
self.trainloader = data.DataLoader(trainset,
batch_size=args.batch_size,
drop_last=True,
shuffle=True,
**kwargs)
self.valloader = data.DataLoader(testset,
batch_size=args.batch_size,
drop_last=False,
shuffle=False,
**kwargs)
self.nclass = trainset.num_class
# model
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=BatchNorm2d,
base_size=args.base_size,
crop_size=args.crop_size,
multi_grid=args.multi_grid,
multi_dilation=args.multi_dilation)
#print(model)
self.logger.info(model)
# optimizer using different LR
if not args.wo_backbone:
params_list = [
{
'params': model.pretrained.parameters(),
'lr': args.lr
},
]
else:
params_list = []
if hasattr(model, 'head'):
params_list.append({
'params': model.head.parameters(),
'lr': args.lr * 10
})
if hasattr(model, 'auxlayer'):
params_list.append({
'params': model.auxlayer.parameters(),
'lr': args.lr * 10
})
optimizer = torch.optim.SGD(params_list,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
self.criterion = SegmentationMultiLosses(nclass=self.nclass)
#self.criterion = SegmentationLosses(se_loss=args.se_loss, aux=args.aux,nclass=self.nclass)
self.model, self.optimizer = model, optimizer
# using cuda
if args.cuda:
self.model = DataParallelModel(self.model).cuda()
self.criterion = DataParallelCriterion(self.criterion).cuda()
# finetune from a trained model
if args.ft:
args.start_epoch = 0
checkpoint = torch.load(args.ft_resume)
if wo_head:
print("WITHout HEAD !!!!!!!!!!")
from collections import OrderedDict
new = OrderedDict()
for k, v in checkpoint['state_dict'].items():
if not k.startswith("head"):
new[k] = v
checkpoint['state_dict'] = new
else:
print("With HEAD !!!!!!!!!!")
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'],
strict=False)
else:
self.model.load_state_dict(checkpoint['state_dict'],
strict=False)
# self.logger.info("=> loaded checkpoint '{}' (epoch {})".format(args.ft_resume, checkpoint['epoch']))
# resuming checkpoint
if args.resume:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
self.logger.info("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# lr scheduler
self.scheduler = utils.LR_Scheduler(args.lr_scheduler,
args.lr,
args.epochs,
len(self.trainloader),
logger=self.logger,
lr_step=args.lr_step)
self.best_pred = 0.0
def test(args):
# output folder
outdir = args.save_folder
if not os.path.exists(outdir):
os.makedirs(outdir)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
testset = get_segmentation_dataset(args.dataset,
split=args.split,
mode=args.mode,
transform=input_transform)
# dataloader
loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
test_data = data.DataLoader(testset,
batch_size=args.test_batch_size,
drop_last=False,
shuffle=False,
collate_fn=test_batchify_fn,
**loader_kwargs)
# model
if args.model_zoo is not None:
model = get_model(args.model_zoo, pretrained=True)
else:
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
dilated=args.dilated,
multi_grid=args.multi_grid,
stride=args.stride,
lateral=args.lateral,
jpu=args.jpu,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=BatchNorm,
base_size=args.base_size,
crop_size=args.crop_size)
# resuming checkpoint
if args.resume is None or not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
# strict=False, so that it is compatible with old pytorch saved models
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# print(model)
scales = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25] if args.dataset == 'citys' else \
[0.5, 0.75, 1.0, 1.25, 1.5, 1.75]
if not args.ms:
scales = [1.0]
evaluator = MultiEvalModule(model,
testset.num_class,
scales=scales,
flip=args.ms).cuda()
evaluator.eval()
metric = utils.SegmentationMetric(testset.num_class)
tbar = tqdm(test_data)
for i, (image, dst) in enumerate(tbar):
if 'val' in args.mode:
with torch.no_grad():
predicts = evaluator.parallel_forward(image)
metric.update(dst, predicts)
pixAcc, mIoU = metric.get()
tbar.set_description('pixAcc: %.4f, mIoU: %.4f' %
(pixAcc, mIoU))
else:
# with torch.no_grad():
# outputs = evaluator.parallel_forward(image)
# predicts = [testset.make_pred(torch.max(output, 1)[1].cpu().numpy())
# for output in outputs]
# for predict, impath in zip(predicts, dst):
# mask = utils.get_mask_pallete(predict, args.dataset)
# outname = os.path.splitext(impath)[0] + '.png'
# mask.save(os.path.join(outdir, outname))
with torch.no_grad():
outputs = evaluator.parallel_forward(image)
# predicts = [testset.make_pred(torch.max(output, 1)[1].cpu().numpy())
# for output in outputs]
predicts = [
torch.softmax(output, 1).cpu().numpy()
for output in outputs
]
for predict, impath in zip(predicts, dst):
# mask = utils.get_mask_pallete(predict, args.dataset)
import numpy as np
from PIL import Image
mask = Image.fromarray(
(predict[0, 1, :, :] * 255).astype(np.uint8))
outname = os.path.splitext(impath)[0] + '.bmp'
mask.save(os.path.join(outdir, outname))
def test(args):
# output folder
outdir = '%s/msdanet_vis' % (args.dataset)
if not os.path.exists(outdir):
os.makedirs(outdir)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
if args.eval:
testset = get_segmentation_dataset(args.dataset,
split='val',
mode='testval',
transform=input_transform)
else: # set split='test' for test set
testset = get_segmentation_dataset(args.dataset,
split='val',
mode='vis',
transform=input_transform)
# dataloader
loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
test_data = data.DataLoader(testset,
batch_size=args.test_batch_size,
drop_last=False,
shuffle=False,
collate_fn=test_batchify_fn,
**loader_kwargs)
if args.model_zoo is not None:
model = get_model(args.model_zoo, pretrained=True)
else:
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=BatchNorm2d,
base_size=args.base_size,
crop_size=args.crop_size,
multi_grid=args.multi_grid,
multi_dilation=args.multi_dilation)
# resuming checkpoint
if args.resume is None or not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
# strict=False, so that it is compatible with old pytorch saved models
model.load_state_dict(checkpoint['state_dict'], strict=False)
print(model)
num_class = testset.num_class
evaluator = MultiEvalModule(model,
testset.num_class,
multi_scales=args.multi_scales).cuda()
evaluator.eval()
tbar = tqdm(test_data)
def eval_batch(image, dst, evaluator, eval_mode):
if eval_mode:
# evaluation mode on validation set
targets = dst
outputs = evaluator.parallel_forward(image)
batch_inter, batch_union, batch_correct, batch_label = 0, 0, 0, 0
for output, target in zip(outputs, targets):
correct, labeled = utils.batch_pix_accuracy(
output.data.cpu(), target)
inter, union = utils.batch_intersection_union(
output.data.cpu(), target, testset.num_class)
batch_correct += correct
batch_label += labeled
batch_inter += inter
batch_union += union
return batch_correct, batch_label, batch_inter, batch_union
else:
# Visualize and dump the results
im_paths = dst
outputs = evaluator.parallel_forward(image)
predicts = [
torch.max(output, 1)[1].cpu().numpy() + testset.pred_offset
for output in outputs
]
for predict, impath in zip(predicts, im_paths):
mask = utils.get_mask_pallete(predict, args.dataset)
outname = os.path.splitext(impath)[0] + '.png'
mask.save(os.path.join(outdir, outname))
# dummy outputs for compatible with eval mode
return 0, 0, 0, 0
total_inter, total_union, total_correct, total_label = \
np.int64(0), np.int64(0), np.int64(0), np.int64(0)
for i, (image, dst) in enumerate(tbar):
if torch_ver == "0.3":
image = Variable(image, volatile=True)
correct, labeled, inter, union = eval_batch(
image, dst, evaluator, args.eval)
else:
with torch.no_grad():
correct, labeled, inter, union = eval_batch(
image, dst, evaluator, args.eval)
pixAcc, mIoU, IoU = 0, 0, 0
if args.eval:
total_correct += correct.astype('int64')
total_label += labeled.astype('int64')
total_inter += inter.astype('int64')
total_union += union.astype('int64')
pixAcc = np.float64(1.0) * total_correct / (
np.spacing(1, dtype=np.float64) + total_label)
IoU = np.float64(1.0) * total_inter / (
np.spacing(1, dtype=np.float64) + total_union)
mIoU = IoU.mean()
tbar.set_description('pixAcc: %.4f, mIoU: %.4f' % (pixAcc, mIoU))
return pixAcc, mIoU, IoU, num_class
sys.path.append(r"E:\Project\PyTorch-Encoding\build\lib.win-amd64-3.6")
import encoding
from encoding.models import get_segmentation_model
from encoding.nn import SyncBatchNorm
from encoding.parallel import DataParallelModel
from osgeo import gdal
import threading
gdal.AllRegister()
# Get the model
checkpoint = torch.load(
r'E:\Project\PyTorch-Encoding\runs\arcs\deeplab\resnest269\model_best.pth.tar\model_best.pth.tar'
)
model = get_segmentation_model("deeplab",
dataset="arcs",
backbone="resnest269",
aux=True,
se_loss=False,
norm_layer=SyncBatchNorm,
base_size=128,
crop_size=128)
model = DataParallelModel(model).cuda()
model.module.load_state_dict(checkpoint['state_dict'])
model.eval()
def processData(tmpName):
oriTileDir = "F:\\色林错\\dataSet\\" + str(tmpName) + r"\OriginTileData"
# maskTileDir = "F:\\色林错\\dataSet\\" + str(tmpName) + r"\MaskTileData"
tmpDir = "F:\\色林错\\dataSet\\" + str(tmpName) + r"\tmpTrainTest"
if not os.path.exists(tmpDir):
os.makedirs(tmpDir)
length = dataSet[tmpName]["length"]
def test(args):
# output folder
outdir = 'outdir'
if not os.path.exists(outdir):
os.makedirs(outdir)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
data_kwargs = {'root': args.data_root}
if args.eval:
testset = get_segmentation_dataset(args.dataset,
split='val',
mode='testval',
transform=input_transform,
**data_kwargs)
elif args.test_val:
testset = get_segmentation_dataset(args.dataset,
split='val',
mode='test',
transform=input_transform,
**data_kwargs)
else:
testset = get_segmentation_dataset(args.dataset,
split='test',
mode='test',
transform=input_transform,
**data_kwargs)
# dataloader
loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
test_data = data.DataLoader(testset,
batch_size=args.test_batch_size,
drop_last=False,
shuffle=False,
collate_fn=test_batchify_fn,
**loader_kwargs)
# model
if args.model_zoo is not None:
model = get_model(args.model_zoo, pretrained=True)
#model.base_size = args.base_size
#model.crop_size = args.crop_size
else:
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=SyncBatchNorm,
base_size=args.base_size,
crop_size=args.crop_size)
# resuming checkpoint
if args.resume is None or not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
# strict=False, so that it is compatible with old pytorch saved models
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
print(model)
# scales = [0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25] if args.dataset == 'citys' else \
# [0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0]
scales = [1.0]
evaluator = MultiEvalModule(model, testset.num_class, scales=scales).cuda()
evaluator.eval()
metric = utils.SegmentationMetric(testset.num_class)
tbar = tqdm(test_data)
for i, (image, dst) in enumerate(tbar):
if args.eval:
with torch.no_grad():
predicts = evaluator.parallel_forward(image)
metric.update(dst, predicts)
pixAcc, mIoU = metric.get()
tbar.set_description('pixAcc: %.4f, mIoU: %.4f' %
(pixAcc, mIoU))
else:
with torch.no_grad():
outputs = evaluator.parallel_forward(image)
predicts = [
testset.make_pred(torch.max(output, 1)[1].cpu().numpy())
for output in outputs
]
for predict, impath in zip(predicts, dst):
mask = utils.get_mask_pallete(predict, args.dataset)
outname = os.path.splitext(impath)[0] + '.png'
mask.save(os.path.join(outdir, outname))
def __init__(self, args):
self.args = args
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size
}
trainset = get_dataset(args.dataset,
split=args.train_split,
mode='train',
**data_kwargs)
testset = get_dataset(args.dataset,
split='val',
mode='val',
**data_kwargs)
self.train_sampler = torch.utils.data.distributed.DistributedSampler(
trainset)
self.val_sampler = torch.utils.data.distributed.DistributedSampler(
testset)
# dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
#self.trainloader = data.DataLoader(trainset, batch_size=args.batch_size,
# drop_last=True, shuffle=True, **kwargs)
#collate_fn=test_batchify_fn,
self.trainloader = data.DataLoader(trainset,
batch_size=args.batch_size //
args.world_size,
drop_last=True,
shuffle=False,
sampler=self.train_sampler,
**kwargs)
#self.valloader = data.DataLoader(testset, batch_size=args.batch_size,
self.valloader = data.DataLoader(testset,
batch_size=args.test_batch_size //
args.world_size,
drop_last=False,
shuffle=False,
sampler=self.val_sampler,
**kwargs)
self.nclass = trainset.num_class
#Norm_method = nn.SyncBatchNorm
#Norm_method = nn.BatchNorm2d(momentum=0.01)
Norm_method = nn.BatchNorm2d
# model
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
multi_grid=args.multi_grid,
se_loss=args.se_loss,
norm_layer=Norm_method,
lateral=args.lateral,
root=args.backbone_path,
base_size=args.base_size,
crop_size=args.crop_size)
if self.args.rank == 0:
print(model)
# optimizer using different LR
params_list = [
{
'params': model.pretrained.parameters(),
'lr': args.lr
},
]
if hasattr(model, 'head'):
params_list.append({
'params': model.head.parameters(),
'lr': args.lr * 10
})
if hasattr(model, 'auxlayer'):
params_list.append({
'params': model.auxlayer.parameters(),
'lr': args.lr * 10
})
optimizer = torch.optim.SGD(params_list,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
self.optimizer = optimizer
#self.model = model
# criterions
self.criterion = SegmentationLosses(se_loss=args.se_loss,
aux=args.aux,
nclass=self.nclass,
se_weight=args.se_weight,
aux_weight=args.aux_weight)
device = torch.device('cuda:{}'.format(args.local_rank))
self.device = device
# using cuda
if args.cuda:
#self.model = DataParallelModel(self.model).cuda()
#self.model = self.model.cuda()
sync_bn_model = FullModel(model, self.criterion)
#self.model.cuda()
#broadcast_params(self.model)
#num_gpus = torch.cuda.device_count()
#local_rank = args.local_rank % num_gpus
#local_rank = args.local_rank
#process_group = torch.distributed.new_group([args.local_rank])
#process_group = torch.distributed.new_group([args.rank])
#sync_bn_model = torch.nn.utils.convert_sync_batchnorm(self.model, process_group)
sync_bn_model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
sync_bn_model)
sync_bn_model = sync_bn_model.to(device)
#self.model = torch.nn.parallel.DistributedDataParallel(self.model, device_ids=[args.local_rank], output_device=args.local_rank)
self.model = torch.nn.parallel.DistributedDataParallel(
sync_bn_model,
device_ids=[args.local_rank],
output_device=args.local_rank,
find_unused_parameters=True)
#self.criterion = DataParallelCriterion(self.criterion).cuda()
#self.criterion = self.criterion.cuda()
dist.barrier()
# resuming checkpoint
#if args.resume is not None and self.args.rank == 0:
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
old_state_dict = checkpoint['state_dict']
new_state_dict = dict()
for k, v in old_state_dict.items():
if k.startswith('module.'):
#new_state_dict[k[len('module.'):]] = old_state_dict[k]
new_state_dict[k] = old_state_dict[k]
else:
new_state_dict[k] = old_state_dict[k]
args.start_epoch = checkpoint['epoch']
if args.cuda:
#self.model.module.load_state_dict(checkpoint['state_dict'])
#self.model.load_state_dict(checkpoint['state_dict'])
self.model.load_state_dict(new_state_dict)
else:
#self.model.load_state_dict(checkpoint['state_dict'])
self.model.load_state_dict(new_state_dict)
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
for state in self.optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.cuda()
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
# lr scheduler
self.scheduler = utils.LR_Scheduler(args.lr_scheduler,
args.lr,
args.epochs,
len(self.trainloader),
local_rank=self.args.rank)
print('len(trainloader) : %.3f ' % (len(self.trainloader)))
self.best_pred = 0.0
#for sumaryWriter
self.track_loss = 0.0
self.track_pixAcc = 0.0
self.track_mIoU = 0.0
def __init__(self, args):
self.args = args
self.logger = utils.create_logger(self.args.exp_dir, "log")
for k, v in vars(self.args).items():
self.logger.info((k, v))
if self.args.cuda:
device = torch.device("cuda")
self.logger.info("training on gpu:" + self.args.gpu_id)
else:
self.logger.info("training on cpu")
device = torch.device("cpu")
self.device = device
#指定随机数
set_seed(args.random_seed)
# args.log_name = str(args.checkname)
#好像是可以加速
cudnn.benchmark = True
#读取数据集，现在只有VOC
self.logger.info('training on dataset ' + self.args.dataset)
input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([.485, .456, .406], [.229, .224, .225])
])
data_kwargs = {
'transform': input_transform,
'root': self.args.data_dir,
'base_size': self.args.base_size,
'crop_size': self.args.crop_size,
'logger': self.logger,
'scale': self.args.scale
}
trainset = get_segmentation_dataset(self.args.dataset,
split='train',
mode='train',
**data_kwargs)
testset = get_segmentation_dataset(self.args.dataset,
split='val',
mode='val',
**data_kwargs)
# dataloader
kwargs = {'num_workers': self.args.num_workers, 'pin_memory': True} \
if self.args.cuda else {}
self.train_iter = data.DataLoader(trainset,
batch_size=self.args.batch_size,
drop_last=True,
shuffle=True,
**kwargs)
self.val_iter = data.DataLoader(testset,
batch_size=self.args.batch_size,
drop_last=False,
shuffle=False,
**kwargs)
self.num_classes = trainset.num_classes
self.input_channels = trainset.input_channels
#create model
kwargs = {'fuse_attention': self.args.fuse_attention}
self.model = get_segmentation_model(args.arch,
dataset=args.dataset,
backbone=args.backbone)
print("=> creating model %s" % self.args.arch)
# self.model = archs.__dict__[self.args.arch](num_classes=self.num_classes,
# input_channels=self.input_channels, **model_kwargs)
self.model = self.model.to(device)
# self.logger.info(self.model)
self.optimizer = None
params = filter(lambda p: p.requires_grad, self.model.parameters())
if self.args.optimizer == "Adam":
self.optimizer = torch.optim.Adam(
params, lr=self.args.lr, weight_decay=self.args.weight_decay)
elif self.args.optimizer == 'SGD':
self.optimizer = torch.optim.SGD(
params,
lr=self.args.lr,
momentum=self.args.momentum,
weight_decay=self.args.weight_decay)
else:
raise NotImplementedError
#loss函数
self.criterion = nn.CrossEntropyLoss(
ignore_index=trainset.IGNORE_INDEX)
#语义分割评价指标
self.metric = SegmentationMetric(self.num_classes)
#学习率策略
self.scheduler = LR_Scheduler(self.args.scheduler, base_lr = self.args.lr, num_epochs=self.args.epochs, \
iters_per_epoch=len(self.train_iter))
#创建实验结果保存目录
self.writer = SummaryWriter(args.exp_dir)
# with open(os.path.join(args.exp_dir,'config.yml'), 'w') as f:
# yaml.dump(config, f)
#用tensoboard看一下模型结构
X, label = next(iter(self.train_iter))
self.writer.add_graph(self.model, X.to(device))
self.epoch_begin = 0
self.best_iou = 0.0
#在训练开始前看看输出是什么
val_log = self.validate(epoch=-1, is_visualize_segmentation=True)
self.write_into_tensorboard(val_log, val_log, epoch=-1)
#checkpoint_PATH
if self.args.checkpoint_PATH is not None:
if self.args.only_read_model:
model, _, _, _, _ = load_checkpoint(model,
self.args.checkpoint_PATH)
else:
model, self.epoch_begin, self.best_iou, self.optimizer = load_checkpoint(
model, self.args.checkpoint_PATH, epoch_begin, best_iou,
optimizer, scheduler)
def test(args):
# output folder
outdir = args.save_folder
if not os.path.exists(outdir):
os.makedirs(outdir)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
testset = get_segmentation_dataset(args.dataset,
split=args.split,
mode=args.mode,
transform=input_transform)
# dataloader
loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
test_data = data.DataLoader(testset,
batch_size=args.test_batch_size,
drop_last=False,
shuffle=False,
collate_fn=test_batchify_fn,
**loader_kwargs)
# model
if args.model_zoo is not None:
model = get_model(args.model_zoo, pretrained=True)
else:
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
dilated=args.dilated,
multi_grid=args.multi_grid,
stride=args.stride,
lateral=args.lateral,
jpu=args.jpu,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=BatchNorm,
base_size=args.base_size,
crop_size=args.crop_size)
# resuming checkpoint
if args.resume is None or not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
# strict=False, so that it is compatible with old pytorch saved models
model.load_state_dict(checkpoint['state_dict'], strict=False)
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# print(model)
scales = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25] if args.dataset == 'citys' else \
[0.5, 0.75, 1.0, 1.25, 1.5, 1.75]
if not args.ms:
scales = [1.0]
evaluator = MultiEvalModule(model,
testset.num_class,
scales=scales,
flip=args.ms).cuda()
evaluator.eval()
tbar = tqdm(test_data)
total_inter, total_union, total_correct, total_label = 0, 0, 0, 0
result = []
for i, (image, dst) in enumerate(tbar):
# print(dst)
with torch.no_grad():
if i > 20:
st = time.time()
outputs = evaluator.forward(image[0].unsqueeze(0).cuda())
if i > 20:
result.append(1 / (time.time() - st))
print(np.mean(result), np.std(result))
if 'val' in args.mode:
# compute image IoU metric
inter, union, area_pred, area_lab = batch_intersection_union(
outputs, dst[0], testset.num_class)
total_label += area_lab
total_inter += inter
total_union += union
class_pixAcc = 1.0 * inter / (np.spacing(1) + area_lab)
class_IoU = 1.0 * inter / (np.spacing(1) + union)
print("img Classes pixAcc:", class_pixAcc)
print("img Classes IoU:", class_IoU)
else:
# save prediction results
predict = testset.make_pred(
torch.max(output, 1)[1].cpu().numpy())
mask = utils.get_mask_pallete(predict, args.dataset)
outname = os.path.splitext(dst[0])[0] + '.png'
mask.save(os.path.join(outdir, outname))
if 'val' in args.mode:
# compute set IoU metric
pixAcc = 1.0 * total_inter / (np.spacing(1) + total_label)
IoU = 1.0 * total_inter / (np.spacing(1) + total_union)
mIoU = IoU.mean()
print("set Classes pixAcc:", pixAcc)
print("set Classes IoU:", IoU)
print("set mean IoU:", mIoU)
def get_model(args):
if args.backbone == "deeplabv2_multi":
model = DeeplabMulti(num_classes=args.num_classes,
pretrained=args.imagenet_pretrained)
elif args.backbone == "new_deeplabv2_multi":
model = NewDeeplabMulti(num_classes=args.num_classes,
pretrained=args.imagenet_pretrained,
use_se=args.use_se,
train_bn=not args.freeze_bn,
norm_style=args.norm_style)
elif args.backbone == 'deeplabv3_resnest50':
from encoding.models import get_segmentation_model
from encoding.nn import SyncBatchNorm
model = get_segmentation_model('deeplab',
dataset='citys',
backbone='resnest50',
aux=True,
norm_layer=SyncBatchNorm)
elif args.backbone == 'deeplabv3_resnest101':
from encoding.models import get_segmentation_model
from encoding.nn import SyncBatchNorm
model = get_segmentation_model('deeplab',
dataset='citys',
backbone='resnest101',
aux=True,
norm_layer=SyncBatchNorm)
elif args.backbone == 'deeplabv2_vgg':
model = DeeplabVGG(num_classes=args.num_classes)
elif args.backbone == 'vgg16_fcn8s':
model = VGG16_FCN8s(num_classes=args.num_classes)
elif args.backbone == 'hrnet':
raise NotImplementedError()
# from graphs.models.hrnet import HighResolutionNet
# model = HighResolutionNet(cfg)
# model.init_weights(self.args.pretrained_ckpt_file)
else:
raise NotImplementedError()
if 'deeplabv2' in args.backbone or 'vgg16_fcn8s' == args.backbone:
params = model.optim_parameters(args)
else:
# https://github.com/zhanghang1989/PyTorch-Encoding/blob/master/experiments/segmentation/train.py#L153
params = [
{
'params': model.pretrained.parameters(),
'lr': args.lr
},
]
if hasattr(model, 'head'):
params.append({
'params': model.head.parameters(),
'lr': args.lr * 10
})
print("Model head has 10x LR")
if hasattr(model, 'auxlayer'):
params.append({
'params': model.auxlayer.parameters(),
'lr': args.lr * 10
})
print("Model auxlayer has 10x LR")
args.numpy_transform = True
return model, params
def test(args):
# output folder
outdir = 'outdir'
if not os.path.exists(outdir):
os.makedirs(outdir)
# data transforms
if not os.path.exists(args.save):
os.makedirs(args.save)
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])])
# dataset
if args.eval:
testset = get_segmentation_dataset(args.dataset, split='val', mode='testval',
transform=input_transform)
else:
testset = get_segmentation_dataset(args.dataset, split='test', mode='test',
transform=input_transform)
# dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
test_data = data.DataLoader(testset, batch_size=args.test_batch_size,
drop_last=False, shuffle=False,
collate_fn=test_batchify_fn, **kwargs)
# model
if args.model_zoo is not None:
model = get_model(args.model_zoo, pretrained=True)
else:
model = get_segmentation_model(args.model, dataset=args.dataset,
backbone = args.backbone, aux = args.aux,
se_loss = args.se_loss, norm_layer = BatchNorm2d)
# resuming checkpoint
if args.resume is None or not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'" .format(args.resume))
checkpoint = torch.load(args.resume)
# strict=False, so that it is compatible with old pytorch saved models
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
print(model)
evaluator = MultiEvalModuleCityscapes(model, testset.num_class).cuda()
evaluator.eval()
tbar = tqdm(test_data)
interp = nn.Upsample(size=(1024, 2048), mode='bilinear')
def eval_batch(image, dst, evaluator, eval_mode, hist, names):
if eval_mode:
# evaluation mode on validation set
targets = dst
outputs = evaluator.parallel_forward(image)
batch_inter, batch_union, batch_correct, batch_label = 0, 0, 0, 0
for output, target, name in zip(outputs, targets, names):
output = interp(output)
correct, labeled = utils.batch_pix_accuracy(output.data.cpu(), target)
inter, union = utils.batch_intersection_union(
output.data.cpu(), target, testset.num_class)
batch_correct += correct
batch_label += labeled
batch_inter += inter
batch_union += union
a = target.numpy().flatten()
b = output.data.cpu()
_, b = torch.max(b, 1)
b = b.numpy().flatten()
n = testset.num_class
k = (a >= 0) & (a < n)
hist += np.bincount(n * a[k].astype(int) + b[k], minlength = n ** 2).reshape(n, n)
output = output.data.cpu().numpy()[0]
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output)
output = Image.fromarray(output)
name = name.split('/')[-1]
output.save('%s/%s' % (args.save, name))
output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0]))
return batch_correct, batch_label, batch_inter, batch_union, hist
else:
# test mode, dump the results
im_paths = dst
outputs = evaluator.parallel_forward(image)
predicts = [torch.max(output, 1)[1].cpu().numpy() + testset.pred_offset
for output in outputs]
for predict, impath in zip(predicts, im_paths):
mask = utils.get_mask_pallete(predict, args.dataset)
outname = os.path.splitext(impath)[0] + '.png'
mask.save(os.path.join(outdir, outname))
# dummy outputs for compatible with eval mode
return 0, 0, 0, 0
total_inter, total_union, total_correct, total_label = \
np.int64(0), np.int64(0), np.int64(0), np.int64(0)
hist =np.zeros((testset.num_class, testset.num_class))
for i, (image, dst, name) in enumerate(tbar):
if torch_ver == "0.3":
image = Variable(image, volatile=True)
correct, labeled, inter, union, hist = eval_batch(image, dst, evaluator, args.eval, hist, name)
else:
with torch.no_grad():
correct, labeled, inter, union, hist = eval_batch(image, dst, evaluator, args.eval, hist, name)
if args.eval:
total_correct += correct.astype('int64')
total_label += labeled.astype('int64')
total_inter += inter.astype('int64')
total_union += union.astype('int64')
pixAcc = np.float64(1.0) * total_correct / (np.spacing(1, dtype=np.float64) + total_label)
IoU = np.float64(1.0) * total_inter / (np.spacing(1, dtype=np.float64) + total_union)
mIoU = IoU.mean()
tbar.set_description(
'pixAcc: %.4f, mIoU: %.4f' % (pixAcc, mIoU))
mIoUs = np.diag(hist) / (hist.sum(1) + hist.sum(0) - np.diag(hist))
print(str(round(np.nanmean(mIoUs) * 100, 2)))
print(mIoUs)
def test(args):
# output folder
outdir = 'outdir'
if not os.path.exists(outdir):
os.makedirs(outdir)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
if args.eval:
testset = get_segmentation_dataset(args.dataset,
split='val',
mode='val',
transform=input_transform)
else:
testset = get_segmentation_dataset(args.dataset,
split='test',
mode='test',
transform=input_transform)
# dataloader
loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
test_data = data.DataLoader(testset,
batch_size=args.test_batch_size,
drop_last=False,
shuffle=False,
collate_fn=test_batchify_fn,
**loader_kwargs)
# model
if args.model_zoo is not None:
model = get_model(args.model_zoo, pretrained=True)
else:
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=BatchNorm2d,
base_size=args.base_size,
crop_size=args.crop_size)
# resuming checkpoint
# if args.resume is None or not os.path.isfile(args.resume):
# raise RuntimeError("=> no checkpoint found at '{}'" .format(args.resume))
# checkpoint = torch.load(args.resume)
# strict=False, so that it is compatible with old pytorch saved models
# pretrained_dict = checkpoint['state_dict']
# model_dict = model.state_dict()
#for name, param in pretrained_dict.items():
# if name not in model_dict:
# continue
# if isinstance(param, Parameter):
# backwards compatibility for serialized parameters
# param = param.data
# model_dict[name].copy_(param)
#model.load_state_dict(checkpoint['state_dict'])
# print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
print(model)
model = model.cuda()
model.eval()
run_time = list()
for i in range(0, 100):
input = torch.randn(1, 3, 512, 512).cuda()
# ensure that context initialization and normal_() operations
# finish before you start measuring time
torch.cuda.synchronize()
torch.cuda.synchronize()
start = time.perf_counter()
with torch.no_grad():
output = model(input)
torch.cuda.synchronize() # wait for mm to finish
end = time.perf_counter()
print(end - start)
run_time.append(end - start)
run_time.pop(0)
print('Mean running time is ', np.mean(run_time))
norm = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
print('begin')
if __name__ == "__main__":
print('please write the parse')
args = Options().parse()
print('i am here')
pretrained = args.resume is None and args.verify is None
if args.model_zoo is not None:
model = get_model(args.model_zoo, pretrained=pretrained)
model.base_size = args.base_size
model.crop_size = args.crop_size
else:
model = get_segmentation_model(args.model, dataset=args.dataset,
backbone=args.backbone, aux = args.aux,
se_loss=args.se_loss,
norm_layer=torch.nn.BatchNorm2d if args.acc_bn else SyncBatchNorm,
base_size=args.base_size, crop_size=args.crop_size)
# model = encoding.models.get_segmentation_model(args.model, dataset=args.dataset, aux=args.aux,
# backbone=args.backbone,
# se_loss=args.se_loss, norm_layer=torch.nn.BatchNorm2d)
print('Creating the model:')
print(model)
model.cuda()
model.eval()
# x = Variable(torch.Tensor(4, 3, 480, 480)).cuda()
rgb = cv2.imread('example.jpg')
rgb = np.transpose(rgb, (2, 0, 1))
def semseg(input_path, output_path=None, with_L0=False):
"""
param:
input_path: str, path of input image
output_path: str, path to save output image
return: tuple, [animal_name, "background"] if pixels of "background" dominate,
["background", animal_name] else.
"""
sys.argv = sys.argv[:1]
option = Options()
args = option.parse()
args.aux = True
args.se_loss = True
args.resume = "./checkpoints/encnet_jpu_res101_pcontext.pth.tar" # model checkpoint
torch.manual_seed(args.seed)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# using L0_smooth to transform the orignal picture
if with_L0:
mid_result = os.path.join(os.path.dirname(input_path), "L0_result.png")
L0_smooth(input_path, mid_result)
input_path = mid_result
# model
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
dilated=args.dilated,
lateral=args.lateral,
jpu=args.jpu,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=BatchNorm,
base_size=args.base_size,
crop_size=args.crop_size)
# resuming checkpoint
if args.resume is None or not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume, map_location=torch.device('cpu'))
# strict=False, so that it is compatible with old pytorch saved models
model.load_state_dict(checkpoint['state_dict'], strict=False)
print("semseg model loaded successfully!")
scales = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25] if args.dataset == 'citys' else \
[0.5, 0.75, 1.0, 1.25, 1.5, 1.75]
if not args.ms:
scales = [1.0]
num_classes = datasets[args.dataset.lower()].NUM_CLASS
evaluator = MultiEvalModule(model,
num_classes,
scales=scales,
flip=args.ms).cuda()
evaluator.eval()
classes = np.array([
'empty', 'aeroplane', 'bag', 'bed', 'bedclothes', 'bench', 'bicycle',
'bird', 'boat', 'book', 'bottle', 'building', 'bus', 'cabinet', 'car',
'cat', 'ceiling', 'chair', 'cloth', 'computer', 'cow', 'cup',
'curtain', 'dog', 'door', 'fence', 'floor', 'flower', 'food', 'grass',
'ground', 'horse', 'keyboard', 'light', 'motorbike', 'mountain',
'mouse', 'person', 'plate', 'platform', 'pottedplant', 'road', 'rock',
'sheep', 'shelves', 'sidewalk', 'sign', 'sky', 'snow', 'sofa', 'table',
'track', 'train', 'tree', 'truck', 'tvmonitor', 'wall', 'water',
'window', 'wood'
])
animals = ['bird', 'cat', 'cow', 'dog', 'horse', 'mouse', 'sheep']
img = input_transform(Image.open(input_path).convert('RGB')).unsqueeze(0)
with torch.no_grad():
output = evaluator.parallel_forward(img)[0]
predict = torch.max(output, 1)[1].cpu().numpy() + 1
pred_idx = np.unique(predict)
pred_label = classes[pred_idx]
print("[SemSeg] ", input_path, ": ", pred_label, sep='')
main_pixels = 0
main_idx = -1
for idx, label in zip(pred_idx, pred_label):
if label in animals:
pixels = np.sum(predict == idx)
if pixels > main_pixels:
main_pixels = pixels
main_idx = idx
background_pixels = np.sum(predict != main_idx)
main_animal = classes[main_idx]
predict[predict != main_idx] = 29
mask_matrix = predict.copy()
if output_path is not None:
mask_matrix[np.where(mask_matrix != 29)] = 1
mask_matrix[np.where(mask_matrix == 29)] = 0
mask = utils.get_mask_pallete(mask_matrix, args.dataset)
mask.save(output_path)
if main_idx < 29:
return predict, (main_animal, "background")
else:
return predict, ("background", main_animal)
def __init__(self, args):
self.args = args
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size
}
trainset = get_dataset(args.dataset,
split=args.train_split,
mode='train',
**data_kwargs)
valset = get_dataset(
args.dataset,
split='val',
mode='ms_val' if args.multi_scale_eval else 'fast_val',
**data_kwargs)
# dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True}
self.trainloader = data.DataLoader(trainset,
batch_size=args.batch_size,
drop_last=True,
shuffle=True,
**kwargs)
if self.args.multi_scale_eval:
kwargs['collate_fn'] = test_batchify_fn
self.valloader = data.DataLoader(valset,
batch_size=args.test_batch_size,
drop_last=False,
shuffle=False,
**kwargs)
self.nclass = trainset.num_class
# model
if args.norm_layer == 'bn':
norm_layer = BatchNorm2d
elif args.norm_layer == 'sync_bn':
assert args.multi_gpu, "SyncBatchNorm can only be used when multi GPUs are available!"
norm_layer = SyncBatchNorm
else:
raise ValueError('Invalid norm_layer {}'.format(args.norm_layer))
model = get_segmentation_model(
args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=norm_layer,
base_size=args.base_size,
crop_size=args.crop_size,
multi_grid=True,
multi_dilation=[2, 4, 8],
only_pam=True,
)
print(model)
# optimizer using different LR
params_list = [
{
'params': model.pretrained.parameters(),
'lr': args.lr
},
]
if hasattr(model, 'head'):
params_list.append({
'params': model.head.parameters(),
'lr': args.lr
})
if hasattr(model, 'auxlayer'):
params_list.append({
'params': model.auxlayer.parameters(),
'lr': args.lr
})
optimizer = torch.optim.SGD(params_list,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# criterions
self.criterion = SegmentationMultiLosses()
self.model, self.optimizer = model, optimizer
# using cuda
if args.multi_gpu:
self.model = DataParallelModel(self.model).cuda()
self.criterion = DataParallelCriterion(self.criterion).cuda()
else:
self.model = self.model.cuda()
self.criterion = self.criterion.cuda()
self.single_device_model = self.model.module if self.args.multi_gpu else self.model
# resuming checkpoint
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
self.single_device_model.load_state_dict(checkpoint['state_dict'])
if not args.ft and not (args.only_val or args.only_vis
or args.only_infer):
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {}), best_pred {}".format(
args.resume, checkpoint['epoch'], checkpoint['best_pred']))
# clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
# lr scheduler
self.lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(
optimizer, 0.6)
self.best_pred = 0.0
def main_worker(gpu, ngpus_per_node, args):
global best_pred
args.gpu = gpu
args.rank = args.rank * ngpus_per_node + gpu
print('rank: {} / {}'.format(args.rank, args.world_size))
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank)
torch.cuda.set_device(args.gpu)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
cudnn.benchmark = True
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size
}
trainset = get_dataset(args.dataset,
split=args.train_split,
mode='train',
**data_kwargs)
valset = get_dataset(args.dataset, split='val', mode='val', **data_kwargs)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
val_sampler = torch.utils.data.distributed.DistributedSampler(
valset, shuffle=False)
# dataloader
loader_kwargs = {
'batch_size': args.batch_size,
'num_workers': args.workers,
'pin_memory': True
}
trainloader = data.DataLoader(trainset,
sampler=train_sampler,
drop_last=True,
**loader_kwargs)
valloader = data.DataLoader(valset, sampler=val_sampler, **loader_kwargs)
nclass = trainset.num_class
# model
model_kwargs = {}
if args.rectify:
model_kwargs['rectified_conv'] = True
model_kwargs['rectify_avg'] = args.rectify_avg
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=DistSyncBatchNorm,
base_size=args.base_size,
crop_size=args.crop_size,
**model_kwargs)
if args.gpu == 0:
print(model)
# optimizer using different LR
params_list = [
{
'params': model.pretrained.parameters(),
'lr': args.lr
},
]
if hasattr(model, 'head'):
params_list.append({
'params': model.head.parameters(),
'lr': args.lr * 10
})
if hasattr(model, 'auxlayer'):
params_list.append({
'params': model.auxlayer.parameters(),
'lr': args.lr * 10
})
optimizer = torch.optim.SGD(params_list,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# criterions
criterion = SegmentationLosses(se_loss=args.se_loss,
aux=args.aux,
nclass=nclass,
se_weight=args.se_weight,
aux_weight=args.aux_weight)
# distributed data parallel
model.cuda(args.gpu)
criterion.cuda(args.gpu)
model = DistributedDataParallel(model, device_ids=[args.gpu])
metric = utils.SegmentationMetric(nclass=nclass)
# resuming checkpoint
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
model.module.load_state_dict(checkpoint['state_dict'])
if not args.ft:
optimizer.load_state_dict(checkpoint['optimizer'])
best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
# lr scheduler
scheduler = utils.LR_Scheduler_Head(args.lr_scheduler, args.lr,
args.epochs, len(trainloader))
def training(epoch):
global best_pred
train_loss = 0.0
model.train()
tic = time.time()
for i, (image, target) in enumerate(trainloader):
scheduler(optimizer, i, epoch, best_pred)
optimizer.zero_grad()
outputs = model(image)
target = target.cuda(args.gpu)
loss = criterion(*outputs, target)
loss.backward()
optimizer.step()
train_loss += loss.item()
if i % 100 == 0 and args.gpu == 0:
iter_per_sec = 100.0 / (
time.time() - tic) if i != 0 else 1.0 / (time.time() - tic)
tic = time.time()
print('Epoch: {}, Iter: {}, Speed: {:.3f} iter/sec, Train loss: {:.3f}'. \
format(epoch, i, iter_per_sec, train_loss / (i + 1)))
def validation(epoch):
# Fast test during the training using single-crop only
global best_pred
is_best = False
model.eval()
metric.reset()
for i, (image, target) in enumerate(valloader):
with torch.no_grad():
#correct, labeled, inter, union = eval_batch(model, image, target)
pred = model(image)[0]
target = target.cuda(args.gpu)
metric.update(target, pred)
pixAcc, mIoU = metric.get()
if i % 100 == 0 and args.gpu == 0:
print('pixAcc: %.3f, mIoU: %.3f' % (pixAcc, mIoU))
if args.gpu == 0:
pixAcc, mIoU = torch_dist_avg(args.gpu, pixAcc, mIoU)
print('pixAcc: %.3f, mIoU: %.3f' % (pixAcc, mIoU))
new_pred = (pixAcc + mIoU) / 2
if new_pred > best_pred:
is_best = True
best_pred = new_pred
utils.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.module.state_dict(),
'optimizer': optimizer.state_dict(),
'best_pred': best_pred,
}, args, is_best)
if args.gpu == 0:
print('Starting Epoch:', args.start_epoch)
print('Total Epoches:', args.epochs)
for epoch in range(args.start_epoch, args.epochs):
tic = time.time()
training(epoch)
if epoch % 10 == 0:
validation(epoch)
elapsed = time.time() - tic
if args.gpu == 0:
print(f'Epoch: {epoch}, Time cost: {elapsed}')
validation(epoch)
def __init__(self, args):
self.args = args
self.log_file = args.resume_dir + '/' + args.log_file
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
data_kwargs = {'transform': input_transform, 'img_size': args.img_size}
# trainset = get_segmentation_dataset(args.dataset, mode='train', augment=True, **data_kwargs)
# valset = get_segmentation_dataset(args.dataset, mode='val', augment=False, **data_kwargs)
testset = get_segmentation_dataset(args.testdata,
mode='test',
augment=False,
whole_image=args.whole_image,
**data_kwargs)
test_vis_set = get_segmentation_dataset(args.testdata,
mode='vis',
augment=False,
whole_image=args.whole_image,
**data_kwargs)
# dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True}
# self.trainloader = data.DataLoader(trainset, batch_size=args.batch_size, drop_last=True, shuffle=True, **kwargs)
# self.valloader = data.DataLoader(valset, batch_size=args.batch_size, drop_last=False, shuffle=False, **kwargs)
self.testloader = data.DataLoader(testset,
batch_size=args.batch_size,
drop_last=False,
shuffle=False,
**kwargs)
self.visloader = data.DataLoader(test_vis_set,
batch_size=args.batch_size,
drop_last=False,
shuffle=False,
**kwargs)
# self.nclass = trainset.NUM_CLASS
# model
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
pretrained=args.pretrained,
aux=args.aux,
is_train=args.is_train,
se_loss=args.se_loss,
batchnorm=SynchronizedBatchNorm2d,
img_size=args.img_size,
dilated=args.dilated,
deep_base=args.deep_base,
multi_grid=args.multi_grid,
multi_dilation=args.multi_dilation,
ensemble=args.ensemble,
resaux=args.resaux,
aggaux=args.aggaux,
output_stride=args.output_stride,
test_size=args.test_img_size,
high_rates=args.high_rates,
aspp_rates=args.aspp_rates,
aspp_out_dim=args.aspp_out_dim,
agg_out_dim=args.agg_out_dim,
up_conv=args.up_conv,
refine=args.refine,
refine_ver=args.refine_ver)
# print(model)
# optimizer using different LR
# params_list = [{'params': model.pretrain_model.parameters(), 'lr': args.lr},
# {'params': model.head.parameters(), 'lr': args.lr * args.head_lr_factor}]
# optimizer = torch.optim.SGD(params_list, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
# optimizer = torch.optim.Adam(params_list, lr=args.lr, betas=(0.9, 0.999), weight_decay=args.weight_decay)
# clear start epoch if fine-turning
# if args.ft:
# args.start_epoch = 0
# criterions
'''
self.criterion = SegmentationLosses(se_loss=args.se_loss, aux=args.aux, nclass=self.nclass,
xentropy_weight=args.xentropy_weight, mae_weight=args.mae_weight,
dice_weight=args.dice_weight, lovasz_weight=args.lovasz_weight, loss1_4=args.loss1_4,
lov1_4=args.lov1_4, margin_weight=args.margin_weight, margin=args.margin,
k=args.k, c=args.c)
self.model, self.optimizer = model, optimizer
'''
self.model = model
# using cuda
# cudnn.benchmark = True
self.model = torch.nn.DataParallel(self.model).cuda()
# self.criterion = self.criterion.cuda()
self.sigmoid = nn.Sigmoid().cuda()
# lr_scheduler
# self.scheduler = utils.LR_Scheduler(args.lr_mode, args.lr, args.epochs, len(self.trainloader))
self.best_pred = 0.0
# resuming chechpoint
if args.resume_dir is not None:
if not os.path.isfile(args.resume_dir + '/checkpoint.pth.tar'):
print('=> no chechpoint found at {}'.format(args.resume_dir))
logger(self.log_file,
'=> no chechpoint found at {}'.format(args.resume_dir))
else:
checkpoint = torch.load(args.resume_dir +
'/checkpoint.pth.tar')
args.start_epoch = checkpoint['epoch']
self.model.module.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print('=> loaded checkpoint {0} (epoch {1})'.format(
args.resume_dir, checkpoint['epoch']))
logger(
self.log_file,
'=> loaded checkpoint {0} (epoch {1})'.format(
args.resume_dir, checkpoint['epoch']))
def main_worker(gpu, ngpus_per_node, args):
global best_pred
args.gpu = gpu
args.rank = args.rank * ngpus_per_node + gpu
print('rank: {} / {}'.format(args.rank, args.world_size))
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank)
torch.cuda.set_device(args.gpu)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
cudnn.benchmark = True
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size
}
trainset = get_dataset(args.dataset,
split=args.train_split,
mode='train',
**data_kwargs)
valset = get_dataset(args.dataset, split='val', mode='val', **data_kwargs)
train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
val_sampler = torch.utils.data.distributed.DistributedSampler(
valset, shuffle=False)
# dataloader
loader_kwargs = {
'batch_size': args.batch_size,
'num_workers': args.workers,
'pin_memory': True
}
trainloader = data.DataLoader(trainset,
sampler=train_sampler,
drop_last=True,
**loader_kwargs)
valloader = data.DataLoader(valset, sampler=val_sampler, **loader_kwargs)
nclass = trainset.num_class
# model
model_kwargs = {}
if args.rectify:
model_kwargs['rectified_conv'] = True
model_kwargs['rectify_avg'] = args.rectify_avg
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=DistSyncBatchNorm,
base_size=args.base_size,
crop_size=args.crop_size,
**model_kwargs)
if args.gpu == 0:
print(model)
# optimizer using different LR
params_list = [
{
'params': model.pretrained.parameters(),
'lr': args.lr
},
]
if hasattr(model, 'head'):
params_list.append({
'params': model.head.parameters(),
'lr': args.lr * 10
})
if hasattr(model, 'auxlayer'):
params_list.append({
'params': model.auxlayer.parameters(),
'lr': args.lr * 10
})
optimizer = torch.optim.SGD(params_list,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optimizer = torch.optim.Adam(params_list,
# lr=args.lr,
# # momentum=args.momentum,
# weight_decay=args.weight_decay)
# criterions
criterion = SegmentationLosses(se_loss=args.se_loss,
aux=args.aux,
nclass=nclass,
se_weight=args.se_weight,
aux_weight=args.aux_weight)
# distributed data parallel
model.cuda(args.gpu)
criterion.cuda(args.gpu)
model = DistributedDataParallel(model, device_ids=[args.gpu])
metric = utils.SegmentationMetric(nclass=nclass)
# resuming checkpoint
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
model.module.load_state_dict(checkpoint['state_dict'])
'''
checkpoint = torch.load(args.resume, map_location='cpu')
args.start_epoch = checkpoint['epoch']
model.module.load_state_dict(checkpoint['state_dict'])
model.cuda()
'''
if not args.ft:
optimizer.load_state_dict(checkpoint['optimizer'])
best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
# lr scheduler
scheduler = utils.LR_Scheduler_Head(args.lr_scheduler, args.lr,
args.epochs, len(trainloader))
# train_losses = [2.855, 2.513, 2.275, 2.128, 2.001, 1.875, 1.855, 1.916, 1.987, 1.915, 1.952]
train_losses = []
def training(epoch):
train_sampler.set_epoch(epoch)
global best_pred
train_loss = 0.0
model.train()
tic = time.time()
for i, (image, target) in enumerate(trainloader):
scheduler(optimizer, i, epoch, best_pred)
optimizer.zero_grad()
outputs = model(image)
target = target.cuda(args.gpu)
loss = criterion(*outputs, target)
loss.backward()
optimizer.step()
train_loss += loss.item()
if i % 100 == 0 and args.gpu == 0:
iter_per_sec = 100.0 / (
time.time() - tic) if i != 0 else 1.0 / (time.time() - tic)
tic = time.time()
print('Epoch: {}, Iter: {}, Speed: {:.3f} iter/sec, Train loss: {:.3f}'. \
format(epoch, i, iter_per_sec, train_loss / (i + 1)))
train_losses.append(train_loss / len(trainloader))
if epoch > 1:
if train_losses[epoch] < train_losses[epoch - 1]:
utils.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.module.state_dict(),
'optimizer': optimizer.state_dict(),
'best_pred': new_preds[(epoch - 1) // 10],
},
args,
False,
filename='checkpoint_train.pth.tar')
plt.plot(train_losses)
plt.xlabel('Epoch')
plt.ylabel('Train_loss')
plt.title('Train_Loss')
plt.grid()
plt.savefig('./loss_fig/train_losses.pdf')
plt.savefig('./loss_fig/train_losses.svg')
plt.close()
# p_m = [(0.3, 0.05), (0.23, 0.54)]
# new_preds = [0.175, 0.392]
p_m = []
new_preds = []
def validation(epoch):
# Fast test during the training using single-crop only
global best_pred
is_best = False
model.eval()
metric.reset()
for i, (image, target) in enumerate(valloader):
with torch.no_grad():
pred = model(image)[0]
target = target.cuda(args.gpu)
metric.update(target, pred)
if i % 100 == 0:
all_metircs = metric.get_all()
all_metircs = utils.torch_dist_sum(args.gpu, *all_metircs)
pixAcc, mIoU = utils.get_pixacc_miou(*all_metircs)
if args.gpu == 0:
print('pixAcc: %.3f, mIoU1: %.3f' % (pixAcc, mIoU))
all_metircs = metric.get_all()
all_metircs = utils.torch_dist_sum(args.gpu, *all_metircs)
pixAcc, mIoU = utils.get_pixacc_miou(*all_metircs)
if args.gpu == 0:
print('pixAcc: %.3f, mIoU2: %.3f' % (pixAcc, mIoU))
p_m.append((pixAcc, mIoU))
plt.plot(p_m)
plt.xlabel('10 Epoch')
plt.ylabel('pixAcc, mIoU')
plt.title('pixAcc, mIoU')
plt.grid()
plt.legend(('pixAcc', 'mIoU'))
plt.savefig('./loss_fig/pixAcc_mIoU.pdf')
plt.savefig('./loss_fig/pixAcc_mIoU.svg')
plt.close()
if args.eval: return
new_pred = (pixAcc + mIoU) / 2
new_preds.append(new_pred)
plt.plot(new_preds)
plt.xlabel('10 Epoch')
plt.ylabel('new_predication')
plt.title('new_predication')
plt.grid()
plt.savefig('./loss_fig/new_predication.pdf')
plt.savefig('./loss_fig/new_predication.svg')
plt.close()
if new_pred > best_pred:
is_best = True
best_pred = new_pred
utils.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.module.state_dict(),
'optimizer': optimizer.state_dict(),
'best_pred': best_pred,
},
args,
is_best,
filename='checkpoint_train_{}.pth.tar'.format(epoch + 1))
if args.export:
if args.gpu == 0:
torch.save(model.module.state_dict(), args.export + '.pth')
return
if args.eval:
validation(args.start_epoch)
return
if args.gpu == 0:
print('Starting Epoch:', args.start_epoch)
print('Total Epoches:', args.epochs)
for epoch in range(args.start_epoch, args.epochs):
tic = time.time()
training(epoch)
if epoch % 10 == 0 or epoch == args.epochs - 1:
validation(epoch)
elapsed = time.time() - tic
if args.gpu == 0:
print(f'Epoch: {epoch}, Time cost: {elapsed}')
validation(epoch)
def __init__(self, args):
self.args = args
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size
}
trainset = get_segmentation_dataset(args.dataset,
split=args.train_split,
mode='train',
**data_kwargs)
testset = get_segmentation_dataset(args.dataset,
split='val',
mode='val',
**data_kwargs)
# dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
self.trainloader = data.DataLoader(trainset,
batch_size=args.batch_size,
drop_last=False,
shuffle=True,
**kwargs)
self.valloader = data.DataLoader(testset,
batch_size=args.batch_size,
drop_last=False,
shuffle=False,
**kwargs)
self.nclass = trainset.num_class
# model
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
dilated=args.dilated,
multi_grid=args.multi_grid,
stride=args.stride,
lateral=args.lateral,
jpu=args.jpu,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=SyncBatchNorm,
base_size=args.base_size,
crop_size=args.crop_size)
# print(model)
# optimizer using different LR
params_list = [
{
'params': model.pretrained.parameters(),
'lr': args.lr
},
]
if hasattr(model, 'jpu') and model.jpu:
params_list.append({
'params': model.jpu.parameters(),
'lr': args.lr * 10
})
if hasattr(model, 'head'):
params_list.append({
'params': model.head.parameters(),
'lr': args.lr * 10
})
if hasattr(model, 'auxlayer'):
params_list.append({
'params': model.auxlayer.parameters(),
'lr': args.lr * 10
})
optimizer = torch.optim.SGD(params_list,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
class_balance_weight = 'None'
if args.dataset == "pcontext60":
class_balance_weight = torch.tensor([
1.3225e-01, 2.0757e+00, 1.8146e+01, 5.5052e+00, 2.2060e+00,
2.8054e+01, 2.0566e+00, 1.8598e+00, 2.4027e+00, 9.3435e+00,
3.5990e+00, 2.7487e-01, 1.4216e+00, 2.4986e+00, 7.7258e-01,
4.9020e-01, 2.9067e+00, 1.2197e+00, 2.2744e+00, 2.0444e+01,
3.0057e+00, 1.8167e+01, 3.7405e+00, 5.6749e-01, 3.2631e+00,
1.5007e+00, 5.5519e-01, 1.0056e+01, 1.8952e+01, 2.6792e-01,
2.7479e-01, 1.8309e+00, 2.0428e+01, 1.4788e+01, 1.4908e+00,
1.9113e+00, 2.6166e+02, 2.3233e-01, 1.9096e+01, 6.7025e+00,
2.8756e+00, 6.8804e-01, 4.4140e+00, 2.5621e+00, 4.4409e+00,
4.3821e+00, 1.3774e+01, 1.9803e-01, 3.6944e+00, 1.0397e+00,
2.0601e+00, 5.5811e+00, 1.3242e+00, 3.0088e-01, 1.7344e+01,
2.1569e+00, 2.7216e-01, 5.8731e-01, 1.9956e+00, 4.4004e+00
])
elif args.dataset == "ade20k":
class_balance_weight = torch.tensor([
0.0772, 0.0431, 0.0631, 0.0766, 0.1095, 0.1399, 0.1502, 0.1702,
0.2958, 0.3400, 0.3738, 0.3749, 0.4059, 0.4266, 0.4524, 0.5725,
0.6145, 0.6240, 0.6709, 0.6517, 0.6591, 0.6818, 0.9203, 0.9965,
1.0272, 1.0967, 1.1202, 1.2354, 1.2900, 1.5038, 1.5160, 1.5172,
1.5036, 2.0746, 2.1426, 2.3159, 2.2792, 2.6468, 2.8038, 2.8777,
2.9525, 2.9051, 3.1050, 3.1785, 3.3533, 3.5300, 3.6120, 3.7006,
3.6790, 3.8057, 3.7604, 3.8043, 3.6610, 3.8268, 4.0644, 4.2698,
4.0163, 4.0272, 4.1626, 4.3702, 4.3144, 4.3612, 4.4389, 4.5612,
5.1537, 4.7653, 4.8421, 4.6813, 5.1037, 5.0729, 5.2657, 5.6153,
5.8240, 5.5360, 5.6373, 6.6972, 6.4561, 6.9555, 7.9239, 7.3265,
7.7501, 7.7900, 8.0528, 8.5415, 8.1316, 8.6557, 9.0550, 9.0081,
9.3262, 9.1391, 9.7237, 9.3775, 9.4592, 9.7883, 10.6705,
10.2113, 10.5845, 10.9667, 10.8754, 10.8274, 11.6427, 11.0687,
10.8417, 11.0287, 12.2030, 12.8830, 12.5082, 13.0703, 13.8410,
12.3264, 12.9048, 12.9664, 12.3523, 13.9830, 13.8105, 14.0345,
15.0054, 13.9801, 14.1048, 13.9025, 13.6179, 17.0577, 15.8351,
17.7102, 17.3153, 19.4640, 17.7629, 19.9093, 16.9529, 19.3016,
17.6671, 19.4525, 20.0794, 18.3574, 19.1219, 19.5089, 19.2417,
20.2534, 20.0332, 21.7496, 21.5427, 20.3008, 21.1942, 22.7051,
23.3359, 22.4300, 20.9934, 26.9073, 31.7362, 30.0784
])
elif args.dataset == "cocostuff":
class_balance_weight = torch.tensor([
4.8557e-02, 6.4709e-02, 3.9255e+00, 9.4797e-01, 1.2703e+00,
1.4151e+00, 7.9733e-01, 8.4903e-01, 1.0751e+00, 2.4001e+00,
8.9736e+00, 5.3036e+00, 6.0410e+00, 9.3285e+00, 1.5952e+00,
3.6090e+00, 9.8772e-01, 1.2319e+00, 1.9194e+00, 2.7624e+00,
2.0548e+00, 1.2058e+00, 3.6424e+00, 2.0789e+00, 1.7851e+00,
6.7138e+00, 2.1315e+00, 6.9813e+00, 1.2679e+02, 2.0357e+00,
2.2933e+01, 2.3198e+01, 1.7439e+01, 4.1294e+01, 7.8678e+00,
4.3444e+01, 6.7543e+01, 1.0066e+01, 6.7520e+00, 1.3174e+01,
3.3499e+00, 6.9737e+00, 2.1482e+00, 1.9428e+01, 1.3240e+01,
1.9218e+01, 7.6836e-01, 2.6041e+00, 6.1822e+00, 1.4070e+00,
4.4074e+00, 5.7792e+00, 1.0321e+01, 4.9922e+00, 6.7408e-01,
3.1554e+00, 1.5832e+00, 8.9685e-01, 1.1686e+00, 2.6487e+00,
6.5354e-01, 2.3801e-01, 1.9536e+00, 1.5862e+00, 1.7797e+00,
2.7385e+01, 1.2419e+01, 3.9287e+00, 7.8897e+00, 7.5737e+00,
1.9758e+00, 8.1962e+01, 3.6922e+00, 2.0039e+00, 2.7333e+00,
5.4717e+00, 3.9048e+00, 1.9184e+01, 2.2689e+00, 2.6091e+02,
4.7366e+01, 2.3844e+00, 8.3310e+00, 1.4857e+01, 6.5076e+00,
2.0854e-01, 1.0425e+00, 1.7386e+00, 1.1973e+01, 5.2862e+00,
1.7341e+00, 8.6124e-01, 9.3702e+00, 2.8545e+00, 6.0123e+00,
1.7560e-01, 1.8128e+00, 1.3784e+00, 1.3699e+00, 2.3728e+00,
6.2819e-01, 1.3097e+00, 4.7892e-01, 1.0268e+01, 1.2307e+00,
5.5662e+00, 1.2867e+00, 1.2745e+00, 4.7505e+00, 8.4029e+00,
1.8679e+00, 1.0519e+01, 1.1240e+00, 1.4975e-01, 2.3146e+00,
4.1265e-01, 2.5896e+00, 1.4537e+00, 4.5575e+00, 7.8143e+00,
1.4603e+01, 2.8812e+00, 1.8868e+00, 7.8131e+01, 1.9323e+00,
7.4980e+00, 1.2446e+01, 2.1856e+00, 3.0973e+00, 4.1270e-01,
4.9016e+01, 7.1001e-01, 7.4035e+00, 2.3395e+00, 2.9207e-01,
2.4156e+00, 3.3211e+00, 2.1300e+00, 2.4533e-01, 1.7081e+00,
4.6621e+00, 2.9199e+00, 1.0407e+01, 7.6207e-01, 2.7806e-01,
3.7711e+00, 1.1852e-01, 8.8280e+00, 3.1700e-01, 6.3765e+01,
6.6032e+00, 5.2177e+00, 4.3596e+00, 6.2965e-01, 1.0207e+00,
1.1731e+01, 2.3935e+00, 9.2767e+00, 1.1023e-01, 3.6947e+00,
1.3943e+00, 2.3407e+00, 1.2112e-01, 2.8518e+00, 2.8195e+00,
1.0078e+00, 1.6614e+00, 6.5307e-01, 1.9070e+01, 2.7231e+00,
6.0769e-01
])
# criterions
self.criterion = SegmentationLosses(se_loss=args.se_loss,
aux=args.aux,
nclass=self.nclass,
se_weight=args.se_weight,
aux_weight=args.aux_weight,
weight=class_balance_weight)
self.model, self.optimizer = model, optimizer
# using cuda
if args.cuda:
self.model = DataParallelModel(self.model).cuda()
self.criterion = DataParallelCriterion(self.criterion).cuda()
# resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
# lr scheduler
self.scheduler = utils.LR_Scheduler(args.lr_scheduler, args.lr,
args.epochs, len(self.trainloader))
from encoding.models import get_model, get_segmentation_model, MultiEvalModule
from lcd import lcdutils
from option import Options
args = Options().parse()
# Get the model
if args.model_zoo is not None:
model = get_model(args.model_zoo, pretrained=True).cuda()
else:
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=BatchNorm2d,
base_size=args.base_size,
crop_size=args.crop_size,
multi_grid=args.multi_grid,
multi_dilation=args.multi_dilation).cuda()
# resuming checkpoint
if args.resume_dir is not None:
args.resume = os.path.join(args.resume_dir, "DANet101.pth.tar")
if args.resume is None or not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
# strict=False, so that it is compatible with old pytorch saved models
model.load_state_dict(checkpoint['state_dict'], strict=False)
model.eval()
def __init__(self, args):
self.args = args
args.log_name = str(args.checkname)
self.logger = utils.create_logger(args.log_root, args.log_name)
# data transforms
input_transform = None
# dataset
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size
}
trainset = get_dataset(args.dataset,
split='train',
mode='train',
**data_kwargs)
testset = get_dataset(args.dataset,
split='val',
mode='val',
**data_kwargs)
# dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
self.trainloader = data.DataLoader(trainset,
batch_size=args.batch_size,
drop_last=True,
shuffle=True,
**kwargs)
self.valloader = data.DataLoader(testset,
batch_size=args.batch_size,
drop_last=False,
shuffle=False,
**kwargs)
self.nclass = trainset.num_class
# model
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=BatchNorm2d,
base_size=args.base_size,
crop_size=args.crop_size)
#print(model)
self.logger.info(model)
# optimizer using different LR
params_list = [
{
'params': model.pretrained.parameters(),
'lr': args.lr
},
]
if hasattr(model, 'head'):
params_list.append({
'params': model.head.parameters(),
'lr': args.lr * 10
})
if hasattr(model, 'auxlayer'):
params_list.append({
'params': model.auxlayer.parameters(),
'lr': args.lr * 10
})
optimizer = torch.optim.SGD(params_list,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
self.criterion = SegmentationMultiLosses(nclass=self.nclass)
#self.criterion = SegmentationLosses(se_loss=args.se_loss, aux=args.aux,nclass=self.nclass)
self.model, self.optimizer = model, optimizer
# using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model).cuda()
self.criterion = torch.nn.DataParallel(self.criterion).cuda()
# finetune from a trained model
if args.ft:
args.start_epoch = 0
checkpoint = torch.load(args.ft_resume)
if args.cuda:
self.model.module.load_state_dict(checkpoint, strict=False)
else:
self.model.load_state_dict(checkpoint, strict=False)
self.logger.info("=> loaded checkpoint '{}' (epoch {})".format(
args.ft_resume, args.start_epoch))
# resuming checkpoint
if args.resume:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
self.logger.info("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# lr scheduler
self.scheduler = utils.LR_Scheduler(
args.lr_scheduler,
args.lr,
args.epochs,
len(self.trainloader), # logger=self.logger,
lr_step=args.lr_step)
self.best_pred = 0.0
def test(args):
# output folder
outdir = 'outdir'
if not os.path.exists(outdir):
os.makedirs(outdir)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
if args.eval:
testset = get_segmentation_dataset(args.dataset,
split='val',
mode='val',
transform=input_transform,
return_file=True)
else:
testset = get_segmentation_dataset(args.dataset,
split='test',
mode='test',
transform=input_transform)
# dataloader
loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
test_data = data.DataLoader(testset,
batch_size=args.test_batch_size,
drop_last=False,
shuffle=False,
collate_fn=test_batchify_fn,
**loader_kwargs)
# model
if args.model_zoo is not None:
model = get_model(args.model_zoo, pretrained=True)
else:
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=BatchNorm2d,
base_size=args.base_size,
crop_size=args.crop_size)
# resuming checkpoint
if args.resume is None or not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
# strict=False, so that it is compatible with old pytorch saved models
pretrained_dict = checkpoint['state_dict']
model_dict = model.state_dict()
for name, param in pretrained_dict.items():
if name not in model_dict:
continue
if isinstance(param, Parameter):
# backwards compatibility for serialized parameters
param = param.data
model_dict[name].copy_(param)
#model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
print(model)
# count parameter number
pytorch_total_params = sum(p.numel() for p in model.parameters())
print("Total number of parameters: %d" % pytorch_total_params)
evaluator = MultiEvalModule(model, testset.num_class).cuda()
evaluator.eval()
tbar = tqdm(test_data)
def eval_batch(image, dst, im_paths, evaluator, eval_mode):
if eval_mode:
# evaluation mode on validation set
targets = dst
outputs = evaluator.parallel_forward(image)
batch_inter, batch_union, batch_correct, batch_label = 0, 0, 0, 0
for output, target in zip(outputs, targets):
correct, labeled = utils.batch_pix_accuracy(
output.data.cpu(), target)
inter, union = utils.batch_intersection_union(
output.data.cpu(), target, testset.num_class)
batch_correct += correct
batch_label += labeled
batch_inter += inter
batch_union += union
# save outputs
predicts = [
torch.max(output, 1)[1].cpu().numpy() # + testset.pred_offset
for output in outputs
]
for predict, impath, target in zip(predicts, im_paths, targets):
mask = utils.get_mask_pallete(predict, args.dataset)
outname = os.path.splitext(impath)[0] + '.png'
mask.save(os.path.join(outdir, outname))
# save ground truth into png format
target = target.data.cpu().numpy()
target = utils.get_mask_pallete(target, args.dataset)
outname = os.path.splitext(impath)[0] + '_gtruth.png'
target.save(os.path.join(outdir, outname))
return batch_correct, batch_label, batch_inter, batch_union
else:
# test mode, dump the results
im_paths = dst
outputs = evaluator.parallel_forward(image)
predicts = [
torch.max(output, 1)[1].cpu().numpy() # + testset.pred_offset
for output in outputs
]
for predict, impath in zip(predicts, im_paths):
mask = utils.get_mask_pallete(predict, args.dataset)
outname = os.path.splitext(impath)[0] + '.png'
mask.save(os.path.join(outdir, outname))
# dummy outputs for compatible with eval mode
return 0, 0, 0, 0
total_inter, total_union, total_correct, total_label = \
np.int64(0), np.int64(0), np.int64(0), np.int64(0)
for i, (image, dst, img_paths) in enumerate(tbar):
if torch_ver == "0.3":
image = Variable(image, volatile=True)
correct, labeled, inter, union = eval_batch(
image, dst, img_paths, evaluator, args.eval)
else:
with torch.no_grad():
correct, labeled, inter, union = eval_batch(
image, dst, img_paths, evaluator, args.eval)
if args.eval:
total_correct += correct.astype('int64')
total_label += labeled.astype('int64')
total_inter += inter.astype('int64')
total_union += union.astype('int64')
pixAcc = np.float64(1.0) * total_correct / (
np.spacing(1, dtype=np.float64) + total_label)
IoU = np.float64(1.0) * total_inter / (
np.spacing(1, dtype=np.float64) + total_union)
mIoU = IoU.mean()
tbar.set_description('pixAcc: %.4f, mIoU: %.4f' % (pixAcc, mIoU))
def test(args):
directory = "runs/val_summary/%s/%s/%s/" % (args.dataset, args.model,
args.resume)
if not os.path.exists(directory):
os.makedirs(directory)
writer = SummaryWriter(directory)
# output folder
outdir = 'outdir'
if not os.path.exists(outdir):
os.makedirs(outdir)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
if args.eval:
testset = get_segmentation_dataset(args.dataset,
split='val',
mode='testval',
transform=input_transform)
elif args.test_val:
testset = get_segmentation_dataset(args.dataset,
split='val',
mode='test',
transform=input_transform)
else:
testset = get_segmentation_dataset(args.dataset,
split='test',
mode='test',
transform=input_transform)
# dataloader
loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
test_data = data.DataLoader(testset,
batch_size=args.test_batch_size,
drop_last=False,
shuffle=False,
collate_fn=test_batchify_fn,
**loader_kwargs)
Norm_method = torch.nn.BatchNorm2d
# model
if args.model_zoo is not None:
model = get_model(args.model_zoo, pretrained=True)
#model.base_size = args.base_size
#model.crop_size = args.crop_size
else:
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
multi_grid=args.multi_grid,
num_center=args.num_center,
norm_layer=Norm_method,
root=args.backbone_path,
base_size=args.base_size,
crop_size=args.crop_size)
# resuming checkpoint
if args.resume is None or not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
# strict=False, so that it is compatible with old pytorch saved models
#model.module.load_state_dict(checkpoint['state_dict'])
old_state_dict = checkpoint['state_dict']
new_state_dict = dict()
for k, v in old_state_dict.items():
if k.startswith('module.'):
#new_state_dict[k[len('module.'):]] = old_state_dict[k]
new_state_dict[k[len('model.module.'):]] = old_state_dict[k]
#new_state_dict[k] = old_state_dict[k]
else:
new_state_dict[k] = old_state_dict[k]
#new_k = 'module.' + k
#new_state_dict[new_k] = old_state_dict[k]
model.load_state_dict(new_state_dict)
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
print(model)
scales = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25] if args.dataset == 'citys' else \
[0.75, 1.0, 1.25, 1.5, 1.75, 2.0]
if args.dataset == 'ade20k':
scales = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0]
if not args.ms:
scales = [1.0]
if args.dataset == 'ade20k':
evaluator = MultiEvalModule2(model,
testset.num_class,
scales=scales,
flip=args.ms).cuda()
else:
evaluator = MultiEvalModule(model,
testset.num_class,
scales=scales,
flip=args.ms).cuda()
evaluator.eval()
metric = utils.SegmentationMetric(testset.num_class)
tbar = tqdm(test_data)
for i, (image, dst) in enumerate(tbar):
if args.eval:
with torch.no_grad():
predicts = evaluator.parallel_forward(image)
metric.update(dst, predicts)
pixAcc, mIoU = metric.get()
tbar.set_description('pixAcc: %.4f, mIoU: %.4f' %
(pixAcc, mIoU))
writer.add_scalar('pixAcc', pixAcc, i)
writer.add_scalar('mIoU', mIoU, i)
else:
with torch.no_grad():
outputs = evaluator.parallel_forward(image)
predicts = [
testset.make_pred(torch.max(output, 1)[1].cpu().numpy())
for output in outputs
]
for predict, impath in zip(predicts, dst):
mask = utils.get_mask_pallete(predict, args.dataset)
outname = os.path.splitext(impath)[0] + '.png'
mask.save(os.path.join(outdir, outname))
writer.close()
def __init__(self, args):
if args.se_loss:
args.checkname = args.checkname + "_se"
self.args = args
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])])
# dataset
data_kwargs = {'transform': input_transform, 'base_size': args.base_size,
'crop_size': args.crop_size}
trainset = get_segmentation_dataset(args.dataset, split='train', mode='train',
**data_kwargs)
testset = get_segmentation_dataset(args.dataset, split='val', mode ='val',
**data_kwargs)
# dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': False} \
if args.cuda else {}
self.trainloader = data.DataLoader(trainset, batch_size=args.batch_size,
drop_last=True, shuffle=True, **kwargs)
self.valloader = data.DataLoader(testset, batch_size=args.batch_size,
drop_last=False, shuffle=False, **kwargs)
self.nclass = trainset.num_class
# model
model = get_segmentation_model(args.model, dataset=args.dataset,
backbone = args.backbone, aux = args.aux,
se_loss = args.se_loss, norm_layer = BatchNorm2d,
base_size=args.base_size, crop_size=args.crop_size)
print(model)
# count parameter number
pytorch_total_params = sum(p.numel() for p in model.parameters())
print("Total number of parameters: %d"%pytorch_total_params)
# optimizer using different LR
params_list = [{'params': model.pretrained.parameters(), 'lr': args.lr},]
if hasattr(model, 'head'):
if args.diflr:
params_list.append({'params': model.head.parameters(), 'lr': args.lr*10})
else:
params_list.append({'params': model.head.parameters(), 'lr': args.lr})
if hasattr(model, 'auxlayer'):
if args.diflr:
params_list.append({'params': model.auxlayer.parameters(), 'lr': args.lr*10})
else:
params_list.append({'params': model.auxlayer.parameters(), 'lr': args.lr})
optimizer = torch.optim.SGD(params_list,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
#optimizer = torch.optim.ASGD(params_list,
# lr=args.lr,
# weight_decay=args.weight_decay)
# criterions
self.criterion = SegmentationLosses(se_loss=args.se_loss, aux=args.aux,
nclass=self.nclass)
self.model, self.optimizer = model, optimizer
# using cuda
if args.cuda:
self.model = DataParallelModel(self.model).cuda()
self.criterion = DataParallelCriterion(self.criterion).cuda()
# resuming checkpoint
if args.resume is not None and len(args.resume)>0:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'" .format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
# load weights for the same model
# self.model.module.load_state_dict(checkpoint['state_dict'])
# model and checkpoint have different strucutures
pretrained_dict = checkpoint['state_dict']
model_dict = self.model.module.state_dict()
for name, param in pretrained_dict.items():
if name not in model_dict:
continue
if isinstance(param, Parameter):
# backwards compatibility for serialized parameters
param = param.data
model_dict[name].copy_(param)
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
# clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
# lr scheduler
self.scheduler = utils.LR_Scheduler(args.lr_scheduler, args.lr,
args.epochs, len(self.trainloader),lr_step=args.lr_step)
self.best_pred = 0.0
def __init__(self, args):
self.args = args
args.log_name = str(args.checkname)
self.logger = utils.create_logger(args.log_root, args.log_name)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size,
'logger': self.logger,
'scale': args.scale
}
trainset = get_segmentation_dataset(args.dataset,
split='train',
mode='train',
**data_kwargs)
testset = get_segmentation_dataset(args.dataset,
split='val',
mode='val',
**data_kwargs)
# dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
self.trainloader = data.DataLoader(trainset,
batch_size=args.batch_size,
drop_last=True,
shuffle=True,
**kwargs)
self.valloader = data.DataLoader(testset,
batch_size=args.batch_size,
drop_last=False,
shuffle=False,
**kwargs)
self.nclass = trainset.num_class
# model
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=BatchNorm2d,
base_size=args.base_size,
crop_size=args.crop_size,
multi_grid=args.multi_grid,
multi_dilation=args.multi_dilation)
#print(model)
self.logger.info(model)
# optimizer using different LR
params_list = [
{
'params': model.pretrained.parameters(),
'lr': args.lr
},
]
if hasattr(model, 'head'):
params_list.append({
'params': model.head.parameters(),
'lr': args.lr * 10
})
if hasattr(model, 'auxlayer'):
params_list.append({
'params': model.auxlayer.parameters(),
'lr': args.lr * 10
})
cityscape_weight = torch.FloatTensor([
0.8373, 0.918, 0.866, 1.0345, 1.0166, 0.9969, 0.9754, 1.0489,
0.8786, 1.0023, 0.9539, 0.9843, 1.1116, 0.9037, 1.0865, 1.0955,
1.0865, 1.1529, 1.0507
])
optimizer = torch.optim.SGD(params_list,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
#weight for class imbalance
# self.criterion = SegmentationMultiLosses(nclass=self.nclass, weight=cityscape_weight)
self.criterion = SegmentationMultiLosses(nclass=self.nclass)
#self.criterion = SegmentationLosses(se_loss=args.se_loss, aux=args.aux,nclass=self.nclass)
self.model, self.optimizer = model, optimizer
# using cuda
if args.cuda:
self.model = DataParallelModel(self.model).cuda()
self.criterion = DataParallelCriterion(self.criterion).cuda()
# finetune from a trained model
if args.ft:
args.start_epoch = 0
checkpoint = torch.load(args.ft_resume)
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'],
strict=False)
else:
self.model.load_state_dict(checkpoint['state_dict'],
strict=False)
self.logger.info("=> loaded checkpoint '{}' (epoch {})".format(
args.ft_resume, checkpoint['epoch']))
# resuming checkpoint
if args.resume:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
self.logger.info("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# lr scheduler
self.scheduler = utils.LR_Scheduler(args.lr_scheduler,
args.lr,
args.epochs,
len(self.trainloader),
logger=self.logger,
lr_step=args.lr_step)
self.best_pred = 0.0
def test(args):
# output folder
outdir = 'outdir'
if not os.path.exists(outdir):
os.makedirs(outdir)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
if args.eval:
testset = get_dataset(args.dataset,
split='val',
mode='testval',
transform=input_transform)
elif args.test_val:
testset = get_dataset(args.dataset,
split='val',
mode='test',
transform=input_transform)
else:
testset = get_dataset(args.dataset,
split='test',
mode='test',
transform=input_transform)
# dataloader
loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
test_data = data.DataLoader(testset,
batch_size=args.test_batch_size,
drop_last=False,
shuffle=False,
collate_fn=test_batchify_fn,
**loader_kwargs)
# model
pretrained = args.resume is None and args.verify is None
if args.model_zoo is not None:
model = get_model(args.model_zoo, pretrained=pretrained)
model.base_size = args.base_size
model.crop_size = args.crop_size
else:
# my
model_kwargs = {}
if args.choice_indices is not None:
assert 'alone_resnest50' in args.backbone
model_kwargs['choice_indices'] = args.choice_indices
#
model = get_segmentation_model(
args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=torch.nn.BatchNorm2d if args.acc_bn else SyncBatchNorm,
base_size=args.base_size,
crop_size=args.crop_size,
**model_kwargs)
# resuming checkpoint
if args.verify is not None and os.path.isfile(args.verify):
print("=> loading checkpoint '{}'".format(args.verify))
model.load_state_dict(torch.load(args.verify, map_location='cpu'))
elif args.resume is not None and os.path.isfile(args.resume):
checkpoint = torch.load(args.resume, map_location='cpu')
# strict=False, so that it is compatible with old pytorch saved models
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}'".format(args.resume))
elif not pretrained:
raise RuntimeError("=> no checkpoint found")
print(model)
if args.acc_bn:
from encoding.utils.precise_bn import update_bn_stats
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size
}
trainset = get_dataset(args.dataset,
split=args.train_split,
mode='train',
**data_kwargs)
trainloader = data.DataLoader(ReturnFirstClosure(trainset),
batch_size=args.batch_size,
drop_last=True,
shuffle=True,
**loader_kwargs)
print('Reseting BN statistics')
#model.apply(reset_bn_statistics)
model.cuda()
update_bn_stats(model, trainloader)
if args.export:
torch.save(model.state_dict(), args.export + '.pth')
return
scales = [0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25] if args.dataset == 'citys' else \
[0.5, 0.75, 1.0, 1.25, 1.5, 1.75]#, 2.0
evaluator = MultiEvalModule(model, testset.num_class, scales=scales).cuda()
evaluator.eval()
metric = utils.SegmentationMetric(testset.num_class)
tbar = tqdm(test_data)
for i, (image, dst) in enumerate(tbar):
if args.eval:
with torch.no_grad():
predicts = evaluator.parallel_forward(image)
metric.update(dst, predicts)
pixAcc, mIoU = metric.get()
tbar.set_description('pixAcc: %.4f, mIoU: %.4f' %
(pixAcc, mIoU))
else:
with torch.no_grad():
outputs = evaluator.parallel_forward(image)
predicts = [
testset.make_pred(torch.max(output, 1)[1].cpu().numpy())
for output in outputs
]
for predict, impath in zip(predicts, dst):
mask = utils.get_mask_pallete(predict, args.dataset)
outname = os.path.splitext(impath)[0] + '.png'
mask.save(os.path.join(outdir, outname))
if args.eval:
print('pixAcc: %.4f, mIoU: %.4f' % (pixAcc, mIoU))
def __init__(self, args):
self.args = args
args.log_name = str(args.checkname)
args.log_root = os.path.join(args.dataset, args.log_root) # dataset/log/
self.logger = utils.create_logger(args.log_root, args.log_name)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])])
# dataset
data_kwargs = {'transform': input_transform, 'base_size': args.base_size,
'crop_size': args.crop_size, 'logger': self.logger,
'scale': args.scale}
trainset = get_segmentation_dataset(args.dataset, split='trainval', mode='trainval',
**data_kwargs)
testset = get_segmentation_dataset(args.dataset, split='val', mode='val', # crop fixed size as model input
**data_kwargs)
# dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
self.trainloader = data.DataLoader(trainset, batch_size=args.batch_size,
drop_last=True, shuffle=True, **kwargs)
self.valloader = data.DataLoader(testset, batch_size=args.batch_size,
drop_last=False, shuffle=False, **kwargs)
self.nclass = trainset.num_class
# model
model = get_segmentation_model(args.model, dataset=args.dataset,
backbone=args.backbone,
norm_layer=BatchNorm2d,
base_size=args.base_size, crop_size=args.crop_size,
)
#print(model)
self.logger.info(model)
# optimizer using different LR
params_list = [{'params': model.pretrained.parameters(), 'lr': args.lr},]
if hasattr(model, 'head'):
print("this model has object, head")
params_list.append({'params': model.head.parameters(), 'lr': args.lr*10})
optimizer = torch.optim.SGD(params_list,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
self.criterion = SegmentationLosses(nclass=self.nclass)
self.model, self.optimizer = model, optimizer
# using cuda
if args.cuda:
self.model = DataParallelModel(self.model).cuda()
self.criterion = DataParallelCriterion(self.criterion).cuda()
# resuming checkpoint
if args.resume:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'" .format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
self.logger.info("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
# lr scheduler
self.scheduler = utils.LR_Scheduler(args.lr_scheduler, args.lr,
args.epochs, len(self.trainloader), logger=self.logger,
lr_step=args.lr_step)
self.best_pred = 0.0
def __init__(self, args):
self.args = args
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size
}
trainset = get_segmentation_dataset(args.dataset,
split=args.train_split,
mode='train',
**data_kwargs)
testset = get_segmentation_dataset(args.dataset,
split='val',
mode='val',
**data_kwargs)
# dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
self.trainloader = data.DataLoader(trainset,
batch_size=args.batch_size,
drop_last=True,
shuffle=True,
**kwargs)
self.valloader = data.DataLoader(testset,
batch_size=args.batch_size,
drop_last=False,
shuffle=False,
**kwargs)
self.nclass = trainset.num_class
# model
model = get_segmentation_model(
args.model,
dataset=args.dataset,
backbone=args.backbone,
dilated=args.dilated,
lateral=args.lateral,
jpu=args.jpu,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=torch.nn.BatchNorm2d, ## BatchNorm2d
base_size=args.base_size,
crop_size=args.crop_size)
print(model)
# optimizer using different LR
params_list = [
{
'params': model.pretrained.parameters(),
'lr': args.lr
},
]
if hasattr(model, 'jpu'):
params_list.append({
'params': model.jpu.parameters(),
'lr': args.lr * 10
})
if hasattr(model, 'head'):
params_list.append({
'params': model.head.parameters(),
'lr': args.lr * 10
})
if hasattr(model, 'auxlayer'):
params_list.append({
'params': model.auxlayer.parameters(),
'lr': args.lr * 10
})
optimizer = torch.optim.SGD(params_list,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# criterions
self.criterion = SegmentationLosses(se_loss=args.se_loss,
aux=args.aux,
nclass=self.nclass,
se_weight=args.se_weight,
aux_weight=args.aux_weight)
self.model, self.optimizer = model, optimizer
# using cuda
if args.cuda:
self.model = DataParallelModel(self.model).cuda()
self.criterion = DataParallelCriterion(self.criterion).cuda()
# resuming checkpoint
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
# lr scheduler
self.scheduler = utils.LR_Scheduler(args.lr_scheduler, args.lr,
args.epochs, len(self.trainloader))
self.best_pred = 0.0
def test(args):
# output folder
outdir = args.save_folder
if not os.path.exists(outdir):
os.makedirs(outdir)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
testset = get_segmentation_dataset(args.dataset,
split=args.split,
mode=args.mode,
transform=input_transform)
# dataloader
loader_kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
test_data = data.DataLoader(testset,
batch_size=args.test_batch_size,
drop_last=False,
shuffle=False,
collate_fn=test_batchify_fn,
**loader_kwargs)
# model
if args.model_zoo is not None:
model = get_model(args.model_zoo, pretrained=True)
else:
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
dilated=args.dilated,
multi_grid=args.multi_grid,
stride=args.stride,
lateral=args.lateral,
jpu=args.jpu,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=BatchNorm,
base_size=args.base_size,
crop_size=args.crop_size)
# resuming checkpoint
if args.resume is None or not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
# strict=False, so that it is compatible with old pytorch saved models
model.load_state_dict(checkpoint['state_dict'], strict=False)
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# print(model)
scales = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25] if args.dataset == 'citys' else \
[0.5, 0.75, 1.0, 1.25, 1.5, 1.75]
if not args.ms:
scales = [1.0]
evaluator = MultiEvalModule(model,
testset.num_class,
scales=scales,
flip=args.ms).cuda()
evaluator.eval()
metric = utils.SegmentationMetric(testset.num_class)
tbar = tqdm(test_data)
total_inter, total_union, total_correct, total_label, all_label = 0, 0, 0, 0, 0
# for i, (image, dst) in enumerate(tbar):
# # print(dst)
# with torch.no_grad():
# outputs = evaluator.parallel_forward(image)[0]
# correct, labeled = batch_pix_accuracy(outputs, dst[0])
# total_correct += correct
# all_label += labeled
# img_pixAcc = 1.0 * correct / (np.spacing(1) + labeled)
# inter, union, area_pred, area_lab = batch_intersection_union(outputs, dst[0], testset.num_class)
# total_label += area_lab
# total_inter += inter
# total_union += union
# class_pixAcc = 1.0 * inter / (np.spacing(1) + area_lab)
# class_IoU = 1.0 * inter / (np.spacing(1) + union)
# class_mIoU = class_IoU.mean()
# print("img pixAcc:", img_pixAcc)
# print("img Classes pixAcc:", class_pixAcc)
# print("img Classes IoU:", class_IoU)
# total_pixAcc = 1.0 * total_correct / (np.spacing(1) + all_label)
# pixAcc = 1.0 * total_inter / (np.spacing(1) + total_label)
# IoU = 1.0 * total_inter / (np.spacing(1) + total_union)
# mIoU = IoU.mean()
# print("set pixAcc:", pixAcc)
# print("set Classes pixAcc:", pixAcc)
# print("set Classes IoU:", IoU)
# print("set mean IoU:", mIoU)
for i, (image, dst) in enumerate(tbar):
if 'val' in args.mode:
with torch.no_grad():
predicts = evaluator.parallel_forward(image)
# metric.update(dst[0], predicts[0])
# pixAcc, mIoU = metric.get()
# tbar.set_description( 'pixAcc: %.4f, mIoU: %.4f' % (pixAcc, mIoU))
else:
with torch.no_grad():
outputs = evaluator.parallel_forward(image)
