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
# 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_1 = get_dataset(
'pascal_voc',
root=os.path.expanduser('/fast/users/a1675776/data/encoding/data'),
split='train',
mode='train',
**data_kwargs)
trainset_2 = get_dataset(
'pascal_aug',
root=os.path.expanduser('/fast/users/a1675776/data/encoding/data'),
split='train',
mode='train',
**data_kwargs)
testset = get_dataset(
'pascal_voc',
root=os.path.expanduser('/fast/users/a1675776/data/encoding/data'),
split='val',
mode='val',
**data_kwargs)
concatenate_trainset = torch.utils.data.ConcatDataset(
[trainset_1, trainset_2])
# dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
self.trainloader = data.DataLoader(concatenate_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_1.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=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, '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.Adam(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,
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 __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 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_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
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
# dataset
<<<<<<< HEAD
# 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)
data_kwargs = {'transform': get_inference_augmentation(), 'base_size': args.base_size}
testset = get_dataset(args.dataset,
root=DATA_DIR,
source=['sample_submission.csv'],
split=['test_1801.npy'],
mode='test',
**data_kwargs)
=======
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)
>>>>>>> upstream/master
# dataloader
def visualize(args):
directory = "runs/%s/%s/%s/vis/results" % (args.dataset, args.model,
args.checkname)
print("visualize directory : ", directory)
if not os.path.exists(directory):
os.mkdir(directory)
# Get the model
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)
model = model.cuda()
# 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'])
# for key in checkpoint['state_dict']:
# print(key)
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
model.eval()
# Prepare the image
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size,
'root': args.data_root
}
testset = get_dataset(args.dataset, split='val', mode='val', **data_kwargs)
testloader = data.DataLoader(testset,
batch_size=1,
drop_last=False,
shuffle=False,
**kwargs)
avg = [.485, .456, .406]
std = [.229, .224, .225]
# visualize
cnt = 0
for i, (image, dst) in enumerate(tqdm(testloader)):
if cnt == 100:
break
prob = np.random.rand(1)[0]
if prob > 0.2:
continue
image = image.cuda()
output = model.evaluate(image)
pred = torch.max(output, 1)[1].cpu().numpy() + 1
mask = encoding.utils.get_mask_pallete(pred, 'pascal_context')
dst = dst.numpy() + 1
gt = encoding.utils.get_mask_pallete(dst, 'pascal_context')
im = image.cpu().numpy().squeeze().transpose(1, 2, 0)
im = im * std + avg
im = im * 255
im = im.astype('uint8')
im = Image.fromarray(im)
target = Image.new('RGB', (480 * 3 + 20, 480), color=(255, 255, 255))
target.paste(im, (0, 0))
target.paste(gt, (490, 0))
target.paste(mask, (980, 0))
target.save('{}/{}.png'.format(directory, str(i)))
cnt += 1
def visualize_attn(args):
directory = "runs/%s/%s/%s/vis/attn" % (args.dataset, args.model,
args.checkname)
print("visualize directory : ", directory)
if not os.path.exists(directory):
os.mkdir(directory)
# Get the model
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)
model = model.cuda()
# print(model)
# print("=================================")
# 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'])
# for key in checkpoint['state_dict']:
# print(key)
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
model.eval()
# Prepare the image
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size,
'root': args.data_root
}
testset = get_dataset(args.dataset, split='val', mode='val', **data_kwargs)
testloader = data.DataLoader(testset,
batch_size=1,
drop_last=False,
shuffle=False,
**kwargs)
avg = [.485, .456, .406]
std = [.229, .224, .225]
# visualize
cnt = 0
for i, (image, dst) in enumerate(tqdm(testloader)):
if cnt == 100:
break
prob = np.random.rand(1)[0]
if prob > 0.2:
continue
image = image.cuda()
output = model.evaluate(image)
cnt += 1
print("{}/100".format(cnt))
self.dice_scores = {'train':[], 'val':[]}
self.iou_scores = {'train':[], 'val':[]}
self.best_loss = float("inf")
# # data transforms
# input_transform = transform.Compose([
# transform.ToTensor(),
# transform.Normalize([.485, .456, .406], [.229, .224, .225])])
# dataset
data_kwargs = {'transform': get_training_augmentation(), 'base_size': args.base_size,
'crop_size': args.crop_size}
<<<<<<< HEAD
trainset = get_dataset(args.dataset,
root=DATA_DIR,
source=['train.csv'],
split=['train_#2fold_11940.npy'],
mode='train',
**data_kwargs)
testset = get_dataset(args.dataset,
root=DATA_DIR,
source=['train.csv'],
split=['valid_#2fold_628.npy'],
mode ='val',
**data_kwargs)
=======
trainset = get_dataset(args.dataset, split=args.train_split, mode='train', **data_kwargs)
testset = get_dataset(args.dataset, split='val', mode ='val', **data_kwargs)
>>>>>>> upstream/master
# dataloader
# kwargs = {'num_workers': args.workers, 'pin_memory': True} \
# if args.cuda else {}
