def __init__(self, args):
"""initialize the Trainer"""
# about gpus
self.cuda = args.cuda
self.gpu_ids = args.gpu_ids
self.num_gpus = len(self.gpu_ids)
self.crf_iter_steps = args.crf_iter_steps
self.output_dir = args.output_dir
self.model = 'test'
# define dataloader
self.val_loader = factory.get_dataset(args.data_dir,
batch_size=1,
dataset=args.dataset,
split=args.train_split)
self.nclass = self.val_loader.NUM_CLASSES
# define network
assert args.seg_model in seg_model_obj_dict.keys()
self.seg_model = args.seg_model
self.seg_model = seg_model_obj_dict[self.seg_model](
num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
norm_layer=torch.nn.BatchNorm2d,
bn_mom=args.bn_mom,
freeze_bn=True)
# define criterion
#self.criterion = torch.nn.CrossEntropyLoss(weight=None, ignore_index=255, reduction='mean')
self.model = full_model.FullModel(seg_model=self.seg_model,
model=self.model)
# define evaluator
self.evaluator = Evaluator(self.nclass)
# using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model,
device_ids=self.gpu_ids)
#patch_replication_callback(self.model)
self.model = self.model.cuda()
#self.criterion = 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))
print('Restore parameters from the {}'.format(args.resume))
checkpoint = torch.load(args.resume)
self.global_step = checkpoint['global_step']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
class Trainer(object):
def __init__(self, args):
"""initialize the Trainer"""
# about gpus
self.cuda = args.cuda
self.gpu_ids = args.gpu_ids
self.num_gpus = len(self.gpu_ids)
self.crf_iter_steps = args.crf_iter_steps
self.output_dir = args.output_dir
# define dataloader
self.val_loader = factory.get_dataset(args.data_dir,
batch_size=1,
dataset=args.dataset,
split=args.train_split)
self.nclass = self.val_loader.NUM_CLASSES
# define network
assert args.seg_model in seg_model_obj_dict.keys()
self.seg_model = args.seg_model
self.model = seg_model_obj_dict[self.seg_model](num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
norm_layer=torch.nn.BatchNorm2d,
bn_mom=args.bn_mom,
freeze_bn=True)
# define evaluator
self.evaluator = Evaluator(self.nclass)
# using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model, device_ids=self.gpu_ids)
#patch_replication_callback(self.model)
self.model = self.model.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))
print('Restore parameters from the {}'.format(args.resume))
checkpoint = torch.load(args.resume)
self.global_step = checkpoint['global_step']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
def validation(self):
"""validation procedure
"""
# set validation mode
self.model.eval()
self.evaluator.reset()
crf_100_steps = 0.0
start = timeit.default_timer()
for i in range(len(self.val_loader)):
#for i, sample in enumerate(self.val_loader):
sample = self.val_loader[i]
image = sample['image']
#image = image.repeat(self.num_gpus, 1, 1, 1)
#print("{}-th sample, Image shape {}, label shape {}".format(i + 1, image.size(), target.size()))
if self.cuda:
image = image.cuda()
image = image.unsqueeze(dim=0)
# forward
with torch.no_grad():
output = self.model(image)
# the output of the pspnet is a tuple
if self.seg_model == 'pspnet':
output = output[0]
# get probs, shape [N, C, H, W] --> [N, H, W, C]
probs = F.softmax(output, dim=1).permute(0, 2, 3, 1).squeeze_()
probs_np = probs.data.cpu().numpy()
#pred = output.data.cpu().numpy()
# CRF post-processing
image_name = self.val_loader.image_lists[i]
#real_image = cv2.cvtColor(cv2.imread(image_name), cv2.COLOR_BGR2RGB)
real_image = Image.open(image_name).convert('RGB')
real_image = np.array(real_image).astype(np.uint8)
crf_start = timeit.default_timer()
pred = crf.dense_crf(real_image=real_image, probs=probs_np, iter_steps=self.crf_iter_steps)
crf_end = timeit.default_timer()
crf_100_steps += (crf_end - crf_start)
# save output
path_to_output = os.path.join(self.output_dir, self.val_loader.image_ids[i] + '.png')
result = Image.fromarray(pred.astype(np.uint8))
result.save(path_to_output)
#cv2.imwrite(path_to_output, pred)
# report time of CRF
if not i % 100:
stop = timeit.default_timer()
print("current step = {} ({:.3f} sec), crf time {:.3f} sec".
format(i, stop - start, crf_100_steps))
crf_100_steps = 0.0
start = timeit.default_timer()
def __init__(self, args):
"""initialize the Trainer"""
# about gpus
self.cuda = args.cuda
self.gpu_ids = args.gpu_ids
self.num_gpus = len(self.gpu_ids)
self.no_val = args.no_val
# about training schedule
self.init_global_step = args.init_global_step
self.start_epoch = args.start_epoch
self.train_epochs = args.epochs
# about the learning rate
self.init_lr = args.init_lr
self.lr_scheduler = args.lr_scheduler
self.slow_start_lr = args.slow_start_lr
self.lr_multiplier = args.lr_multiplier
self.accumulation_steps = args.accumulation_steps
# about the model_dir and checkpoint
self.model_dir = args.model_dir
self.save_ckpt_steps = args.save_ckpt_steps
self.max_ckpt_nums = args.max_ckpt_nums
self.saved_ckpt_filenames = []
self.checkname = args.checkname
# define global setp
self.global_step = 0
self.main_gpu = args.main_gpu
# sync bn, both can be used.
self.norm_layer = syncbn.BatchNorm2d if args.sync_bn else nn.BatchNorm2d
#self.norm_layer = SynchronizedBatchNorm2d if args.sync_bn else nn.BatchNorm2d
# define tensorboard summary
self.train_writer = SummaryWriter(log_dir=self.model_dir)
self.val_writer = SummaryWriter(
log_dir=os.path.join(self.model_dir, 'eval'))
# define dataloader
self.train_loader, self.nclass = factory.get_data_loader(
args.data_dir,
batch_size=args.batch_size,
crop_size=args.crop_size,
dataset=args.dataset,
split=args.train_split,
num_workers=args.workers,
pin_memory=True)
self.val_loader, _ = factory.get_data_loader(
args.data_dir,
dataset=args.dataset,
split="test" if 'camvid' in args.dataset else "val")
# max iters
self.steps_per_epochs = len(self.train_loader)
self.max_iter = self.steps_per_epochs * self.train_epochs
# define network
assert args.seg_model in seg_model_obj_dict.keys()
self.seg_model = args.seg_model
self.model = seg_model_obj_dict[self.seg_model](
num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
norm_layer=self.norm_layer,
bn_mom=args.bn_mom,
freeze_bn=args.freeze_bn)
# define criterion
self.loss_type = args.loss_type
self.criterion = loss_factory.criterion_choose(
self.nclass,
weight=None,
loss_type=args.loss_type,
ignore_index=255,
reduction='mean',
max_iter=self.max_iter,
args=args)
self.model_with_loss = full_model.FullModel(seg_model=self.seg_model,
model=self.model,
loss_type=self.loss_type,
criterion=self.criterion)
# define evaluator
self.evaluator = Evaluator(self.nclass)
# using cuda
# If you need to move a model to GPU via .cuda(), please do so before constructing optimizers for it.
# Parameters of a model after .cuda() will be different objects with those before the call.
# In general, you should make sure that optimized parameters
# live in consistent locations when optimizers are constructed and used.
if args.cuda:
self.model_with_loss = torch.nn.DataParallel(
self.model_with_loss, device_ids=self.gpu_ids)
if self.norm_layer is SynchronizedBatchNorm2d:
patch_replication_callback(self.model_with_loss)
print("INFO:PyTorch: The batch norm layer is {}".format(
self.norm_layer))
elif self.norm_layer is syncbn.BatchNorm2d:
parallel.patch_replication_callback(self.model_with_loss)
print("INFO:PyTorch: The batch norm layer is Hang Zhang's {}".
format(self.norm_layer))
self.model_with_loss = self.model_with_loss.cuda(self.main_gpu)
# optimizer parameters, construct optim after module
self.params_list = []
self.params_list = model_init.seg_model_get_optim_params(
self.params_list,
self.model_with_loss.module.model,
norm_layer=self.norm_layer,
seg_model=args.seg_model,
base_lr=args.init_lr,
lr_multiplier=self.lr_multiplier,
weight_decay=args.weight_decay)
self.optimizer = torch.optim.SGD(self.params_list,
momentum=args.momentum,
nesterov=args.nesterov)
# define learning rate scheduler.
# Be careful about the learning rate for different params list, check
# the `params_list` and the `lr_scheduler` to ensure the strategy is right.
self.scheduler = train_utils.lr_scheduler(
init_lr=self.init_lr,
mode=self.lr_scheduler,
num_epochs=self.train_epochs,
max_iter=self.max_iter,
slow_start_steps=args.slow_start_steps,
slow_start_lr=args.slow_start_lr,
multiplier=self.lr_multiplier)
# resuming checkpoint
#self.best_pred = 0.0
if args.resume is not None:
if os.path.isfile(args.resume):
#raise RuntimeError("=> no checkpoint found at '{}'" .format(args.resume))
print("INFO:PyTorch: Restore checkpoint from {}".format(
args.resume))
checkpoint = torch.load(args.resume)
self.global_step = checkpoint['global_step']
if args.cuda:
self.model_with_loss.module.load_state_dict(
checkpoint['state_dict'])
else:
self.model_with_loss.load_state_dict(
checkpoint['state_dict'])
self.start_epoch = (self.global_step +
1) // self.steps_per_epochs
class Trainer(object):
def __init__(self, args):
"""initialize the Trainer"""
# about gpus
self.cuda = args.cuda
self.gpu_ids = args.gpu_ids
self.num_gpus = len(self.gpu_ids)
self.no_val = args.no_val
# about training schedule
self.init_global_step = args.init_global_step
self.start_epoch = args.start_epoch
self.train_epochs = args.epochs
# about the learning rate
self.init_lr = args.init_lr
self.lr_scheduler = args.lr_scheduler
self.slow_start_lr = args.slow_start_lr
self.lr_multiplier = args.lr_multiplier
self.accumulation_steps = args.accumulation_steps
# about the model_dir and checkpoint
self.model_dir = args.model_dir
self.save_ckpt_steps = args.save_ckpt_steps
self.max_ckpt_nums = args.max_ckpt_nums
self.saved_ckpt_filenames = []
self.checkname = args.checkname
# define global setp
self.global_step = 0
self.main_gpu = args.main_gpu
# sync bn, both can be used.
self.norm_layer = syncbn.BatchNorm2d if args.sync_bn else nn.BatchNorm2d
#self.norm_layer = SynchronizedBatchNorm2d if args.sync_bn else nn.BatchNorm2d
# define tensorboard summary
self.train_writer = SummaryWriter(log_dir=self.model_dir)
self.val_writer = SummaryWriter(
log_dir=os.path.join(self.model_dir, 'eval'))
# define dataloader
self.train_loader, self.nclass = factory.get_data_loader(
args.data_dir,
batch_size=args.batch_size,
crop_size=args.crop_size,
dataset=args.dataset,
split=args.train_split,
num_workers=args.workers,
pin_memory=True)
self.val_loader, _ = factory.get_data_loader(
args.data_dir,
dataset=args.dataset,
split="test" if 'camvid' in args.dataset else "val")
# max iters
self.steps_per_epochs = len(self.train_loader)
self.max_iter = self.steps_per_epochs * self.train_epochs
# define network
assert args.seg_model in seg_model_obj_dict.keys()
self.seg_model = args.seg_model
self.model = seg_model_obj_dict[self.seg_model](
num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
norm_layer=self.norm_layer,
bn_mom=args.bn_mom,
freeze_bn=args.freeze_bn)
# define criterion
self.loss_type = args.loss_type
self.criterion = loss_factory.criterion_choose(
self.nclass,
weight=None,
loss_type=args.loss_type,
ignore_index=255,
reduction='mean',
max_iter=self.max_iter,
args=args)
self.model_with_loss = full_model.FullModel(seg_model=self.seg_model,
model=self.model,
loss_type=self.loss_type,
criterion=self.criterion)
# define evaluator
self.evaluator = Evaluator(self.nclass)
# using cuda
# If you need to move a model to GPU via .cuda(), please do so before constructing optimizers for it.
# Parameters of a model after .cuda() will be different objects with those before the call.
# In general, you should make sure that optimized parameters
# live in consistent locations when optimizers are constructed and used.
if args.cuda:
self.model_with_loss = torch.nn.DataParallel(
self.model_with_loss, device_ids=self.gpu_ids)
if self.norm_layer is SynchronizedBatchNorm2d:
patch_replication_callback(self.model_with_loss)
print("INFO:PyTorch: The batch norm layer is {}".format(
self.norm_layer))
elif self.norm_layer is syncbn.BatchNorm2d:
parallel.patch_replication_callback(self.model_with_loss)
print("INFO:PyTorch: The batch norm layer is Hang Zhang's {}".
format(self.norm_layer))
self.model_with_loss = self.model_with_loss.cuda(self.main_gpu)
# optimizer parameters, construct optim after module
self.params_list = []
self.params_list = model_init.seg_model_get_optim_params(
self.params_list,
self.model_with_loss.module.model,
norm_layer=self.norm_layer,
seg_model=args.seg_model,
base_lr=args.init_lr,
lr_multiplier=self.lr_multiplier,
weight_decay=args.weight_decay)
self.optimizer = torch.optim.SGD(self.params_list,
momentum=args.momentum,
nesterov=args.nesterov)
# define learning rate scheduler.
# Be careful about the learning rate for different params list, check
# the `params_list` and the `lr_scheduler` to ensure the strategy is right.
self.scheduler = train_utils.lr_scheduler(
init_lr=self.init_lr,
mode=self.lr_scheduler,
num_epochs=self.train_epochs,
max_iter=self.max_iter,
slow_start_steps=args.slow_start_steps,
slow_start_lr=args.slow_start_lr,
multiplier=self.lr_multiplier)
# resuming checkpoint
#self.best_pred = 0.0
if args.resume is not None:
if os.path.isfile(args.resume):
#raise RuntimeError("=> no checkpoint found at '{}'" .format(args.resume))
print("INFO:PyTorch: Restore checkpoint from {}".format(
args.resume))
checkpoint = torch.load(args.resume)
self.global_step = checkpoint['global_step']
if args.cuda:
self.model_with_loss.module.load_state_dict(
checkpoint['state_dict'])
else:
self.model_with_loss.load_state_dict(
checkpoint['state_dict'])
self.start_epoch = (self.global_step +
1) // self.steps_per_epochs
# 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
def training(self, epoch):
"""training procedure
"""
# set training mode
self.model_with_loss.train()
self.evaluator.reset()
start_time = time.time()
self.optimizer.zero_grad()
# training loop
for i, sample in enumerate(self.train_loader):
# set grad zero
self.optimizer.zero_grad()
# accumulate global steps
if (i + 1) % self.accumulation_steps == 0:
self.global_step += 1
image, target = sample['image'], sample['label']
if self.cuda:
image, target = image.cuda(self.main_gpu), target.cuda(
self.main_gpu)
# adjust learning rate, pass input through the model, update
self.scheduler(self.optimizer, self.global_step, epoch)
output, loss = self.model_with_loss(inputs=image,
target=target,
global_step=self.global_step)
#print(target.size())
loss = loss.mean()
loss.backward()
#if (i + 1) % self.accumulation_steps == 0:
# update the parameters and set the gradient to 0.
self.optimizer.step()
# add batch sample into evaluator
pred = np.argmax(output.data.cpu().numpy(), axis=1)
target = target.cpu().numpy()
self.evaluator.add_batch(target, pred)
# log info per 20 steps
#if self.global_step % 20 == 0 and (i + 1) % self.accumulation_steps == 0:
if self.global_step % 20 == 0:
# the used time
used_time = time.time() - start_time
px_acc = self.evaluator.pixel_accuracy_np()
miou = self.evaluator.mean_iou_np()
lr_now = self.optimizer.param_groups[0]['lr']
print(
"INFO:PyTorch: epoch={}/{}, steps={}, loss={:.5f}, learning_rate={:.5f}, train_miou={:.5f}, px_accuracy={:.5f}"
" ({:.3f} sec)".format(epoch + 1,
self.train_epochs, self.global_step,
loss.item(), lr_now, miou, px_acc,
used_time))
# summary per 100 steps
if self.global_step % 100 == 0:
self.train_writer.add_scalar('train_miou',
miou,
global_step=self.global_step)
self.train_writer.add_scalar('px_accuracy',
px_acc,
global_step=self.global_step)
self.train_writer.add_scalar('learning_rate',
lr_now,
global_step=self.global_step)
self.train_writer.add_scalar('train_loss',
loss.item(),
global_step=self.global_step)
start_time = time.time()
# save checkpoints
if (self.global_step % self.save_ckpt_steps
== 0) or (i == self.steps_per_epochs - 1):
filename = os.path.join(
self.model_dir,
"{0}_ckpt_{1}.pth".format(self.checkname,
self.global_step))
self.saved_ckpt_filenames.append(filename)
# remove the newest file if saved ckpts if more than max_ckpt_nums
if len(self.saved_ckpt_filenames) > self.max_ckpt_nums:
del_filename = self.saved_ckpt_filenames.pop(0)
os.remove(del_filename)
# save new ckpt
state = {
'global_step': self.global_step,
'state_dict': self.model_with_loss.module.state_dict(),
'optimizer': self.optimizer.state_dict(),
}
torch.save(state, filename)
def validation(self, epoch):
"""validation procedure
"""
# set validation mode
self.model_with_loss.eval()
self.evaluator.reset()
test_loss = 0.0
for i, sample in enumerate(self.val_loader):
image, target = sample['image'], sample['label']
# repeat, you can uncomment this the line
#image, target = image.repeat(self.num_gpus, 1, 1, 1), target.repeat(self.num_gpus, 1, 1)
if self.cuda:
image, target = image.cuda(self.main_gpu), target.cuda(
self.main_gpu)
# forward
with torch.no_grad():
output = self.model_with_loss(inputs=image, mode='val')
# Add batch sample into evaluator
pred = np.argmax(output.data.cpu().numpy(), axis=1)
target = target.cpu().numpy()
self.evaluator.add_batch(target, pred)
# log and summary the validation results
px_acc = self.evaluator.pixel_accuracy_np()
val_miou = self.evaluator.mean_iou_np(is_show_per_class=True)
print(
"\nINFO:PyTorch: validation results: miou={:5f}, px_acc={:5f}, loss={:5f} \n"
.format(val_miou, px_acc, test_loss))
self.val_writer.add_scalar('val_loss', test_loss, self.global_step)
self.val_writer.add_scalar('val_miou', val_miou, self.global_step)
self.val_writer.add_scalar('val_px_acc', px_acc, self.global_step)
class Trainer(object):
def __init__(self, args):
"""initialize the Trainer"""
# about gpus
self.cuda = args.cuda
self.gpu_ids = args.gpu_ids
self.num_gpus = len(self.gpu_ids)
self.crf_iter_steps = args.crf_iter_steps
self.output_dir = args.output_dir
self.model = 'val'
# define dataloader
self.val_loader = factory.get_dataset(args.data_dir,
batch_size=1,
dataset=args.dataset,
split=args.train_split)
self.nclass = self.val_loader.NUM_CLASSES
# define network
assert args.seg_model in seg_model_obj_dict.keys()
self.seg_model = args.seg_model
self.seg_model = seg_model_obj_dict[self.seg_model](num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
norm_layer=torch.nn.BatchNorm2d,
bn_mom=args.bn_mom,
freeze_bn=True)
# define criterion
self.criterion = torch.nn.CrossEntropyLoss(weight=None, ignore_index=255, reduction='mean')
self.model = full_model.FullModel(seg_model=self.seg_model,
model=self.model,
criterion=self.criterion)
# define evaluator
self.evaluator = Evaluator(self.nclass)
# using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model, device_ids=self.gpu_ids)
#patch_replication_callback(self.model)
self.model = self.model.cuda()
self.criterion = 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))
print('Restore parameters from the {}'.format(args.resume))
checkpoint = torch.load(args.resume)
self.global_step = checkpoint['global_step']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
def validation(self):
"""validation procedure
"""
# set validation mode
self.model.eval()
self.evaluator.reset()
test_loss = 0.0
start = timeit.default_timer()
for i in range(len(self.val_loader)):
#for i, sample in enumerate(self.val_loader):
sample = self.val_loader[i]
image, target = sample['image'], sample['label']
image, target = image.repeat(self.num_gpus, 1, 1, 1), target.repeat(self.num_gpus, 1, 1)
#print("{}-th sample, Image shape {}, label shape {}".format(i + 1, image.size(), target.size()))
if self.cuda:
image, target = image.cuda(), target.cuda()
# forward
with torch.no_grad():
output = self.model(image)
# the output of the pspnet is a tuple
if self.seg_model == 'pspnet':
output = output[0]
loss = self.criterion(output, target.long())
test_loss += loss.item()
# get probs, shape [N, C, H, W] --> [N, H, W, C]
output = output.squeeze_()
pred = output.data.cpu().numpy()
pred = np.argmax(pred, axis=0)
target = target.squeeze_().cpu().numpy()
# save output
color_img = True
path_to_output = os.path.join(self.output_dir, self.val_loader.image_ids[i] + '.png')
pred = pred.astype(np.uint8)
if color_img:
pass
pred_color = utils.decode_segmap(pred, dataset='pascal')
result = Image.fromarray(pred_color.astype(np.uint8))
result.save(path_to_output)
else:
result = Image.fromarray()
result.save(path_to_output)
# report time
if not i % 100:
stop = timeit.default_timer()
print("current step = {} ({:.3f} sec)".format(i, stop - start))
start = timeit.default_timer()
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# log and summary the validation results
# log and summary the validation results
px_acc = self.evaluator.pixel_accuracy_np()
val_miou = self.evaluator.mean_iou_np(is_show_per_class=True)
print("\nINFO:PyTorch: validation results: miou={:5f}, px_acc={:5f}, loss={:5f} \n".
format(val_miou, px_acc, test_loss))
class Trainer(object):
def __init__(self, args):
"""initialize the Trainer"""
# about gpus
self.cuda = args.cuda
self.gpu_ids = args.gpu_ids
self.num_gpus = len(self.gpu_ids)
self.crf_iter_steps = args.crf_iter_steps
self.output_dir = args.output_dir
self.model = 'test'
# define dataloader
self.val_loader = factory.get_dataset(args.data_dir,
batch_size=1,
dataset=args.dataset,
split=args.train_split)
self.nclass = self.val_loader.NUM_CLASSES
# define network
assert args.seg_model in seg_model_obj_dict.keys()
self.seg_model = args.seg_model
self.seg_model = seg_model_obj_dict[self.seg_model](
num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
norm_layer=torch.nn.BatchNorm2d,
bn_mom=args.bn_mom,
freeze_bn=True)
# define criterion
#self.criterion = torch.nn.CrossEntropyLoss(weight=None, ignore_index=255, reduction='mean')
self.model = full_model.FullModel(seg_model=self.seg_model,
model=self.model)
# define evaluator
self.evaluator = Evaluator(self.nclass)
# using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model,
device_ids=self.gpu_ids)
#patch_replication_callback(self.model)
self.model = self.model.cuda()
#self.criterion = 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))
print('Restore parameters from the {}'.format(args.resume))
checkpoint = torch.load(args.resume)
self.global_step = checkpoint['global_step']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
def validation(self):
"""validation procedure
"""
# set validation mode
self.model.eval()
self.evaluator.reset()
start = timeit.default_timer()
for i in range(len(self.val_loader)):
sample = self.val_loader[i]
image = sample['image']
if self.cuda:
image = image.cuda()
image = image.unsqueeze(dim=0)
# forward
with torch.no_grad():
output = self.model(image)
# the output of the pspnet is a tuple
if self.seg_model == 'pspnet':
output = output[0]
output = output.squeeze_()
pred = output.data.cpu().numpy()
# save output
pred = np.argmax(pred, axis=0)
path_to_output = os.path.join(
self.output_dir, self.val_loader.image_ids[i] + '.png')
result = Image.fromarray(pred.astype(np.uint8))
result.save(path_to_output)
#cv2.imwrite(path_to_output, pred)
# report time of CRF
if not i % 100:
stop = timeit.default_timer()
print("current step = {} ({:.3f} sec)".format(i, stop - start))
start = timeit.default_timer()