def validation(self, epoch):
def _get_pred(batch_im):
with torch.no_grad():
# metric.update also accepts list, so no need to gather results from multi gpus
if self.args.multi_scale_eval:
assert len(batch_im) <= torch.cuda.device_count(
), "Multi-scale testing only allows batch size <= number of GPUs"
scattered_pred = self.ms_evaluator.parallel_forward(
batch_im)
else:
outputs = self.model(batch_im)
scattered_pred = [
out[0] for out in outputs
] if self.args.multi_gpu else [outputs[0]]
return scattered_pred
# Lazy creation
if not hasattr(self, 'ms_evaluator'):
self.ms_evaluator = MultiEvalModule(self.single_device_model,
self.nclass,
scales=self.args.eval_scales,
crop=self.args.crop_eval)
self.metric = utils.SegmentationMetric(self.nclass)
self.model.eval()
tbar = tqdm(self.valloader, desc='\r')
for i, (batch_im, target) in enumerate(tbar):
# No need to put target to GPU, since the metrics are calculated by numpy.
# And no need to put data to GPU manually if we use data parallel.
if not self.args.multi_gpu and not isinstance(
batch_im, (list, tuple)):
batch_im = batch_im.cuda()
scattered_pred = _get_pred(batch_im)
scattered_target = []
ind = 0
for p in scattered_pred:
target_tmp = target[ind:ind + len(p)]
# Multi-scale testing. In fact, len(target_tmp) == 1
if isinstance(target_tmp, (list, tuple)):
assert len(target_tmp) == 1
target_tmp = torch.stack(target_tmp)
scattered_target.append(target_tmp)
ind += len(p)
self.metric.update(scattered_target, scattered_pred)
pixAcc, mIoU = self.metric.get()
tbar.set_description('ep {}, pixAcc: {:.4f}, mIoU: {:.4f}'.format(
epoch + 1, pixAcc, mIoU))
return self.metric.get()
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 = '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 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 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 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 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 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)
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):
