def train(model, loader, optimizer, scheduler=None):
print('now LR: ' + str(optimizer.param_groups[0]['lr']))
model.train()
total, correct, train_loss = 0.0, 0.0, 0.0
for step, data in enumerate(loader):
x, y = data
x, y = x.to(device), y.to(device)
out = model(x)
# print("size:", x.size(), y.size(), out.size())
loss = nn.CrossEntropyLoss()(out, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if scheduler:
scheduler.step()
_, pred = torch.max(out.data, 1)
total += y.size(0)
correct += (pred == y).squeeze().sum().cpu().numpy()
train_loss += loss.item()
if step % 100 == 0:
print("step: {0}, loss:{1}".format(step, loss.item()))
train_acc = correct / total
train_loss /= step
logx.msg("train_acc:" + str(train_acc))
return train_acc, train_loss
def load_weights(net, optimizer, snapshot_file, restore_optimizer_bool=False):
"""
Load weights from snapshot file
"""
logx.msg("Loading weights from model {}".format(snapshot_file))
net, optimizer = restore_snapshot(net, optimizer, snapshot_file, restore_optimizer_bool)
return net, optimizer
def find_cityscapes_images(self, cities, img_root, mask_root, img_ext,
mask_ext, fine_coarse='gtFine'):
"""
Find image and segmentation mask files and return a list of
tuples of them.
Inputs:
img_root: path to parent directory of train/val/test dirs
mask_root: path to parent directory of train/val/test dirs
img_ext: image file extension
mask_ext: mask file extension
cities: a list of cities, each element in the form of 'train/a_city'
or 'val/a_city', for example.
"""
items = []
for city in cities:
img_dir = '{root}/{city}'.format(root=img_root, city=city)
for file_name in os.listdir(img_dir):
basename, ext = os.path.splitext(file_name)
assert ext == '.' + img_ext, '{} {}'.format(ext, img_ext)
full_img_fn = os.path.join(img_dir, file_name)
basename, ext = file_name.split('_leftImg8bit_to_darknight')
if cfg.DATASET.CUSTOM_COARSE_PROB and fine_coarse != 'gtFine':
mask_fn = f'{basename}_leftImg8bit.png'
cc_path = cfg.DATASET.CITYSCAPES_CUSTOMCOARSE
full_mask_fn = os.path.join(cc_path, city, mask_fn)
os.path.isfile(full_mask_fn)
else:
mask_fn = f'{basename}_{fine_coarse}_labelIds{mask_ext}'
full_mask_fn = os.path.join(mask_root, city, mask_fn)
items.append((full_img_fn, full_mask_fn))
logx.msg('mode {} found {} images'.format(self.mode, len(items)))
return items
def init_weights(self, pretrained=cfg.MODEL.HRNET_CHECKPOINT):
# logx.msg('=> init weights from normal distribution')
for name, m in self.named_modules():
if any(part in name for part in {'cls', 'aux', 'ocr'}):
# print('skipped', name)
continue
if isinstance(m, nn.Conv2d):
nn.init.normal_(m.weight, std=0.001)
elif isinstance(m, cfg.MODEL.BNFUNC):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
if os.path.isfile(pretrained):
pretrained_dict = torch.load(pretrained,
map_location={'cuda:0': 'cpu'})
logx.msg('=> loading pretrained model {}'.format(pretrained))
model_dict = self.state_dict()
pretrained_dict = {
k.replace('last_layer', 'aux_head').replace('model.', ''): v
for k, v in pretrained_dict.items()
}
#print(set(model_dict) - set(pretrained_dict))
#print(set(pretrained_dict) - set(model_dict))
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict.keys()
}
model_dict.update(pretrained_dict)
self.load_state_dict(model_dict)
elif pretrained:
raise RuntimeError('No such file {}'.format(pretrained))
def print_evaluate_results(hist,
iu,
epoch=0,
iou_per_scale=None,
log_multiscale_tb=False):
"""
If single scale:
just print results for default scale
else
print all scale results
Inputs:
hist = histogram for default scale
iu = IOU for default scale
iou_per_scale = iou for all scales
"""
id2cat = cfg.DATASET_INST.trainid_to_name
# id2cat = {i: i for i in range(cfg.DATASET.NUM_CLASSES)}
iu_FP = hist.sum(axis=1) - np.diag(hist)
iu_FN = hist.sum(axis=0) - np.diag(hist)
iu_TP = np.diag(hist)
logx.msg('IoU:')
header = ['Id', 'label']
header.extend(['iU_{}'.format(scale) for scale in iou_per_scale])
header.extend(['TP', 'FP', 'FN', 'Precision', 'Recall'])
tabulate_data = []
for class_id in range(len(iu)):
class_data = []
class_data.append(class_id)
class_name = "{}".format(
id2cat[class_id]) if class_id in id2cat else ''
class_data.append(class_name)
for scale in iou_per_scale:
class_data.append(iou_per_scale[scale][class_id] * 100)
total_pixels = hist.sum()
class_data.append(100 * iu_TP[class_id] / total_pixels)
class_data.append(iu_FP[class_id] / iu_TP[class_id])
class_data.append(iu_FN[class_id] / iu_TP[class_id])
class_data.append(iu_TP[class_id] /
(iu_TP[class_id] + iu_FP[class_id]))
class_data.append(iu_TP[class_id] /
(iu_TP[class_id] + iu_FN[class_id]))
tabulate_data.append(class_data)
if log_multiscale_tb:
logx.add_scalar("xscale_%0.1f/%s" % (0.5, str(id2cat[class_id])),
float(iou_per_scale[0.5][class_id] * 100), epoch)
logx.add_scalar("xscale_%0.1f/%s" % (1.0, str(id2cat[class_id])),
float(iou_per_scale[1.0][class_id] * 100), epoch)
logx.add_scalar("xscale_%0.1f/%s" % (2.0, str(id2cat[class_id])),
float(iou_per_scale[2.0][class_id] * 100), epoch)
print_str = str(tabulate((tabulate_data), headers=header, floatfmt='1.2f'))
logx.msg(print_str)
def fit(self, num_eval_per_epoch: int = 10):
steps_per_eval = len(self.train_dataloader) // num_eval_per_epoch
self.train(steps_per_eval)
test_dataloader = self.get_test_dataloader(self.input_dir, self.batch_size)
mean_loss, metrics_scores, _, _ = self.validate(test_dataloader)
logx.msg("Scores on test set: ")
logx.msg(str(metrics_scores))
def train(train_loader, model, criterion, optimizer, epoch, args):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to train mode
model.train()
end = time.time()
for i, (input, target) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
if args.gpu is not None:
input = input.cuda(args.gpu, non_blocking=True)
target = target.cuda(args.gpu, non_blocking=True)
# compute output
output = model(input)
loss = criterion(output, target)
# measure accuracy and record loss
acc1, acc5 = accuracy(output, target, topk=(1, 5))
losses.update(loss.item(), input.size(0))
top1.update(acc1[0], input.size(0))
top5.update(acc5[0], input.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
metrics = {
'loss': losses.avg,
'top1': float(top1.avg),
'top5': float(top5.avg)
}
logx.metric('train', metrics, i + epoch * len(train_loader))
if i % args.print_freq == 0:
logx.msg('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Acc@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
epoch,
i,
len(train_loader),
batch_time=batch_time,
data_time=data_time,
loss=losses,
top1=top1,
top5=top5))
def __init__(self, input_dir, output_dir, model, device,
per_gpu_batch_size, n_gpu, batch_size, learning_rate,
weight_decay, n_epoch, seed, top_k, **kwargs):
# construct param dict
self.construct_param_dict = OrderedDict({
"input_dir": str(input_dir),
"output_dir": str(output_dir),
"learning_rate": learning_rate,
"n_epoch": n_epoch,
"per_gpu_batch_size": per_gpu_batch_size,
"weight_decay": weight_decay,
"seed": seed,
"top_k": top_k,
})
# build log
logx.initialize(logdir=output_dir,
coolname=True,
tensorboard=True,
no_timestamp=False,
hparams={
"solver_construct_dict": self.construct_param_dict,
"model_construct_dict": model.model_construct_dict
},
eager_flush=True)
# arguments
self.record_training_loss_per_epoch = kwargs.pop(
"record_training_loss_per_epoch", False)
self.input_dir = input_dir
self.output_dir = output_dir
self.top_k = top_k
# training utilities
self.model = model
# data utilities
self.train_dataloader = kwargs.pop("train_dataloader", None)
self.dev_dataloader = kwargs.pop("dev_dataloader", None)
self.batch_size = batch_size
self.n_epoch = n_epoch
self.seed = seed
# device
self.device = device
self.n_gpu = n_gpu
logx.msg(f'Number of GPU: {self.n_gpu}.')
self.criterion = kl_div_add_mse_loss
# optimizer and scheduler
if self.train_dataloader:
self.optimizer, self.scheduler = self.get_optimizer(
named_parameters=self.model.named_parameters(),
learning_rate=learning_rate,
weight_decay=weight_decay,
train_dataloader=self.train_dataloader,
n_epoch=n_epoch)
# set up random seeds and model location
self.setup()
def __init__(self, input_dir: Path, output_dir: Path, model: nn.Module,
device: torch.device, per_gpu_batch_size: int, n_gpu: int,
batch_size: int, learning_rate: float, weight_decay: float,
n_epoch: int, seed: int, **kwargs):
# construct param dict
self.construct_param_dict = OrderedDict({
"input_dir": str(input_dir),
"output_dir": str(output_dir),
"learning_rate": learning_rate,
"n_epoch": n_epoch,
"per_gpu_batch_size": per_gpu_batch_size,
"weight_decay": weight_decay,
"seed": seed
})
self.subreddit_embedding_device = torch.device("cuda")
# build log
logx.initialize(logdir=output_dir,
coolname=True,
tensorboard=True,
no_timestamp=False,
hparams={
"solver_hparams": self.state_dict(),
"model_hparams": model.param_dict()
},
eager_flush=True)
# arguments
self.input_dir = input_dir
self.output_dir = output_dir
# training utilities
self.model = model
# data utilities
self.train_dataloader = kwargs.pop("train_dataloader", None)
self.dev_dataloader = kwargs.pop("dev_dataloader", None)
self.batch_size = batch_size
self.n_epoch = n_epoch
self.seed = seed
# device
self.device = device
self.n_gpu = n_gpu
logx.msg(f'Number of GPU: {self.n_gpu}.')
self.criterion = nn.MSELoss()
# optimizer and scheduler
if self.train_dataloader:
self.optimizer, self.scheduler = self.get_optimizer(
named_parameters=self.model.named_parameters(),
learning_rate=learning_rate,
weight_decay=weight_decay,
train_dataloader=self.train_dataloader,
n_epoch=n_epoch)
# set up random seeds and model location
self.setup()
def __init__(self,
weight=None,
ignore_index=cfg.DATASET.IGNORE_LABEL,
reduction="mean"):
super(CrossEntropyLoss2d, self).__init__()
logx.msg("Using Cross Entropy Loss")
self.nll_loss = nn.NLLLoss(weight,
reduction=reduction,
ignore_index=ignore_index)
def __init__(self, mode, quality='fine', joint_transform_list=None,
img_transform=None, label_transform=None, eval_folder=None):
super(Loader, self).__init__(quality=quality, mode=mode,
joint_transform_list=joint_transform_list,
img_transform=img_transform,
label_transform=label_transform)
######################################################################
# Cityscapes-specific stuff:
######################################################################
self.root = cfg.DATASET.CITYSCAPES_DARK_DIR
self.id_to_trainid = cityscapes_labels.label2trainid
self.trainid_to_name = cityscapes_labels.trainId2name
self.fill_colormap()
img_ext = 'jpg'
mask_ext = '.png'
img_root = path.join(self.root, 'leftImg8bit')
mask_root = path.join(self.root, 'gtFine')
if mode == 'folder':
self.all_imgs = make_dataset_folder(eval_folder)
else:
self.fine_cities = cities_cv_split(self.root, mode, cfg.DATASET.CV)
self.all_imgs = self.find_cityscapes_images(
self.fine_cities, img_root, mask_root, img_ext, mask_ext)
logx.msg(f'cn num_classes {self.num_classes}')
self.fine_centroids = uniform.build_centroids(self.all_imgs,
self.num_classes,
self.train,
cv=cfg.DATASET.CV,
id2trainid=self.id_to_trainid)
self.centroids = self.fine_centroids
if cfg.DATASET.COARSE_BOOST_CLASSES and mode == 'train':
self.coarse_cities = coarse_cities(self.root)
img_root = path.join(self.root,
'leftImg8bit_trainextra/leftImg8bit')
mask_root = path.join(self.root, 'gtCoarse', 'gtCoarse')
self.coarse_imgs = self.find_cityscapes_images(
self.coarse_cities, img_root, mask_root, img_ext, mask_ext,
fine_coarse='gtCoarse')
if cfg.DATASET.CLASS_UNIFORM_PCT:
custom_coarse = (cfg.DATASET.CUSTOM_COARSE_PROB is not None)
self.coarse_centroids = uniform.build_centroids(
self.coarse_imgs, self.num_classes, self.train,
coarse=(not custom_coarse), custom_coarse=custom_coarse,
id2trainid=self.id_to_trainid)
for cid in cfg.DATASET.COARSE_BOOST_CLASSES:
self.centroids[cid].extend(self.coarse_centroids[cid])
else:
self.all_imgs.extend(self.coarse_imgs)
self.build_epoch()
def torch_version_float():
version_str = torch.__version__
version_re = re.search(r"^([0-9]+\.[0-9]+)", version_str)
if version_re:
version = float(version_re.group(1))
logx.msg(f"Torch version: {version}, {version_str}")
else:
version = 1.0
logx.msg(f"Can't parse torch version ({version}), assuming {version}")
return version
def __init__(self, classes, weight=None, size_average=False,
ignore_index=cfg.DATASET.IGNORE_LABEL,
norm=False, upper_bound=1.0):
super(EdgeWeightedCrossEntropyLoss2d, self).__init__()
logx.msg("Using Per Image based weighted loss")
self.num_classes = classes
self.nll_loss = nn.NLLLoss2d(weight, size_average,ignore_index)
self.norm = norm
self.upper_bound = upper_bound
self.batch_weights = cfg.BATCH_WEIGHTING
def _load_pretrained_model(self):
pretrained_model = cfg.MODEL.X71_CHECKPOINT
ckpt = torch.load(pretrained_model, map_location='cpu')
model_dict = {k.replace('module.', ''): v for k, v in
ckpt['model_dict'].items()}
state_dict = self.state_dict()
state_dict.update(model_dict)
self.load_state_dict(state_dict, strict=False)
del ckpt
logx.msg('Loaded {} weights'.format(pretrained_model))
def coarse_cities(root):
"""
Find coarse cities
"""
split = 'train_extra'
coarse_path = path.join(root, 'leftImg8bit_trainextra/leftImg8bit', split)
coarse_cities = [f'{split}/' + c for c in os.listdir(coarse_path)]
logx.msg(f'found {len(coarse_cities)} coarse cities')
return coarse_cities
def torch_version_float():
version_str = torch.__version__
version_re = re.search(r'^([0-9]+\.[0-9]+)', version_str)
if version_re:
version = float(version_re.group(1))
logx.msg(f'Torch version: {version}, {version_str}')
else:
version = 1.0
logx.msg(f'Can\'t parse torch version ({version}), assuming {version}')
return version
def __init__(self, classes, weight=None, ignore_index=cfg.DATASET.IGNORE_LABEL,
norm=False, upper_bound=1.0, fp16=False):
super(ImageBasedCrossEntropyLoss2d, self).__init__()
logx.msg("Using Per Image based weighted loss")
self.num_classes = classes
self.nll_loss = nn.NLLLoss(weight, reduction='mean',
ignore_index=ignore_index)
self.norm = norm
self.upper_bound = upper_bound
self.batch_weights = cfg.BATCH_WEIGHTING
self.fp16 = fp16
def get_net(args, criterion):
"""
Get Network Architecture based on arguments provided
"""
net = get_model(network='network.' + args.arch,
num_classes=cfg.DATASET.NUM_CLASSES,
criterion=criterion)
num_params = sum([param.nelement() for param in net.parameters()])
logx.msg('Model params = {:2.1f}M'.format(num_params / 1000000))
net = net.cuda()
return net
def get_trunk(trunk_name, output_stride=8):
"""
Retrieve the network trunk and channel counts.
"""
assert output_stride == 8, 'Only stride8 supported right now'
if trunk_name == 'wrn38':
#
# FIXME: pass in output_stride once we support stride 16
#
backbone = wrn38(pretrained=True)
s2_ch = 128
s4_ch = 256
high_level_ch = 4096
elif trunk_name == 'xception71':
backbone = xception71(output_stride=output_stride,
BatchNorm=Norm2d,
pretrained=True)
s2_ch = 64
s4_ch = 128
high_level_ch = 2048
elif trunk_name == 'seresnext-50' or trunk_name == 'seresnext-101':
backbone = get_resnet(trunk_name, output_stride=output_stride)
s2_ch = 48
s4_ch = -1
high_level_ch = 2048
elif trunk_name == 'resnet-50' or trunk_name == 'resnet-101':
backbone = get_resnet(trunk_name, output_stride=output_stride)
s2_ch = 256
s4_ch = -1
high_level_ch = 2048
elif trunk_name == 'hrnetv2':
backbone = hrnetv2.get_seg_model()
high_level_ch = backbone.high_level_ch
s2_ch = -1
s4_ch = -1
elif trunk_name == 'mobilenetv3_large':
backbone = MobileNetV3_Large(pretrained=False)
s2_ch = 16
s4_ch = 24
high_level_ch = 320
elif trunk_name == 'mobilenetv3_small':
backbone = MobileNetV3_Small(pretrained=False)
s2_ch = 16
s4_ch = 16
high_level_ch = 256
else:
raise 'unknown backbone {}'.format(trunk_name)
logx.msg("Trunk: {}".format(trunk_name))
return backbone, s2_ch, s4_ch, high_level_ch
def __init__(self,
mode,
quality='semantic',
joint_transform_list=None,
img_transform=None,
label_transform=None,
eval_folder=None):
super(Loader, self).__init__(quality=quality,
mode=mode,
joint_transform_list=joint_transform_list,
img_transform=img_transform,
label_transform=label_transform)
root = cfg.DATASET.CITYSURFACES_DIR
self.id_to_trainid = label2trainid
self.trainid_to_name = trainId2name
self.fill_colormap()
######################################################################
# Assemble image lists
######################################################################
if mode == 'folder':
img_ext = 'png'
mask_ext = 'png'
img_root = os.path.join(root, 'val', 'images')
mask_root = os.path.join(root, 'val', 'annotations')
self.all_imgs = self.find_images(img_root, mask_root, img_ext,
mask_ext)
elif mode == 'test':
self.all_imgs = make_dataset_folder(eval_folder, test_mode=True)
else:
splits = {'train': 'train', 'val': 'val', 'test': 'tests'}
split_name = splits[mode]
img_ext = 'png'
mask_ext = 'png'
img_root = os.path.join(root, split_name, 'images')
mask_root = os.path.join(root, split_name, 'annotations')
self.all_imgs = self.find_images(img_root, mask_root, img_ext,
mask_ext)
logx.msg('all imgs {}'.format(len(self.all_imgs)))
self.fine_centroids = uniform.build_centroids(
self.all_imgs,
self.num_classes,
self.train,
cv=cfg.DATASET.CV,
id2trainid=self.id_to_trainid)
self.centroids = self.fine_centroids
self.build_epoch()
def validate(val_loader, model, criterion, args, epoch):
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to evaluate mode
model.eval()
with torch.no_grad():
end = time.time()
for i, (input, target) in enumerate(val_loader):
if args.gpu is not None:
input = input.cuda(args.gpu, non_blocking=True)
target = target.cuda(args.gpu, non_blocking=True)
# compute output
output = model(input)
loss = criterion(output, target)
# measure accuracy and record loss
acc1, acc5 = accuracy(output, target, topk=(1, 5))
losses.update(loss.item(), input.size(0))
top1.update(acc1[0], input.size(0))
top5.update(acc5[0], input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0:
logx.msg('Test: [{0}/{1}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t'
'Acc@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
i,
len(val_loader),
batch_time=batch_time,
loss=losses,
top1=top1,
top5=top5))
logx.msg(' * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'.format(
top1=top1, top5=top5))
metrics = {'top1': float(top1.avg), 'top5': float(top5.avg)}
logx.metric('val', metrics, epoch)
return top1.avg
def cities_cv_split(root, split, cv_split):
"""
Find cities that correspond to a given split of the data. We split the data
such that a given city belongs to either train or val, but never both. cv0
is defined to be the default split.
all_cities = [x x x x x x x x x x x x]
val:
split0 [x x x ]
split1 [ x x x ]
split2 [ x x x ]
trn:
split0 [ x x x x x x x x x]
split1 [x x x x x x x x x]
split2 [x x x x x x x x ]
split - train/val/test
cv_split - 0,1,2,3
cv_split == 3 means use train + val
"""
trn_path = path.join(root, 'leftImg8bit_trainvaltest/leftImg8bit', 'train')
val_path = path.join(root, 'leftImg8bit_trainvaltest/leftImg8bit', 'val')
trn_cities = ['train/' + c for c in os.listdir(trn_path)]
trn_cities = sorted(trn_cities) # sort to insure reproducibility
val_cities = ['val/' + c for c in os.listdir(val_path)]
all_cities = val_cities + trn_cities
if cv_split == 3:
logx.msg('cv split {} {} {}'.format(split, cv_split, all_cities))
return all_cities
num_val_cities = len(val_cities)
num_cities = len(all_cities)
offset = cv_split * num_cities // cfg.DATASET.CV_SPLITS
cities = []
for j in range(num_cities):
if j >= offset and j < (offset + num_val_cities):
if split == 'val':
cities.append(all_cities[j])
else:
if split == 'train':
cities.append(all_cities[j])
logx.msg('cv split {} {} {}'.format(split, cv_split, cities))
return cities
def restore_snapshot(net, optimizer, snapshot, restore_optimizer_bool):
"""
Restore weights and optimizer (if needed ) for resuming job.
"""
checkpoint = torch.load(snapshot, map_location=torch.device('cpu'))
logx.msg("Checkpoint Load Compelete")
if optimizer is not None and 'optimizer' in checkpoint and restore_optimizer_bool:
optimizer.load_state_dict(checkpoint['optimizer'])
if 'state_dict' in checkpoint:
net = forgiving_state_restore(net, checkpoint['state_dict'])
else:
net = forgiving_state_restore(net, checkpoint)
return net, optimizer
def valid(model, loader):
model.eval()
total, correct, valid_loss = 0.0, 0.0, 0.0
with torch.no_grad():
for step, data in enumerate(loader):
x, y = data
x, y = x.to(device), y.to(device)
out = model(x)
loss = nn.CrossEntropyLoss()(out, y)
_, pred = torch.max(out.data, 1)
total += y.size(0)
correct += (pred == y).squeeze().sum().cpu().numpy()
valid_loss += loss.item()
valid_acc = correct / total
valid_loss /= step
logx.msg("valid_acc:" + str(valid_acc))
return valid_acc, valid_loss
def forgiving_state_restore(net, loaded_dict):
"""
Handle partial loading when some tensors don't match up in size.
Because we want to use models that were trained off a different
number of classes.
"""
net_state_dict = net.state_dict()
new_loaded_dict = {}
for k in net_state_dict:
new_k = k
if new_k in loaded_dict and net_state_dict[k].size() == loaded_dict[new_k].size():
new_loaded_dict[k] = loaded_dict[new_k]
else:
logx.msg("Skipped loading parameter {}".format(k))
net_state_dict.update(new_loaded_dict)
net.load_state_dict(net_state_dict)
return net
def train(args, model, device, train_loader, optimizer, epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
if batch_idx % args.log_interval == 0:
logx.msg('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
# capture metrics
metrics = {'loss': loss.item()}
iteration = epoch * len(train_loader) + batch_idx
logx.metric('train', metrics, iteration)
def __init__(self,
mode,
quality='fine',
joint_transform_list=None,
img_transform=None,
label_transform=None,
eval_folder=None):
super(Loader, self).__init__(quality=quality,
mode=mode,
joint_transform_list=joint_transform_list,
img_transform=img_transform,
label_transform=label_transform)
root = cfg.DATASET.KITTI_DIR
######################################################################
# Assemble image lists
######################################################################
if mode == 'folder':
self.all_imgs = make_dataset_folder(eval_folder)
else:
splits = {
'train': 'training',
'val': 'training',
'test': 'training'
}
split_name = splits[mode]
img_ext = 'png'
mask_ext = 'png'
img_root = os.path.join(root, split_name, 'image_2')
mask_root = os.path.join(root, split_name, 'semantic_rgb')
self.all_imgs = self.find_images(img_root, mask_root, img_ext,
mask_ext)
import ipdb
ipdb.set_trace()
self.fill_colormap()
logx.msg('all imgs {}'.format(len(self.all_imgs)))
self.centroids = uniform.build_centroids(self.all_imgs,
self.num_classes,
self.train,
cv=cfg.DATASET.CV)
self.build_epoch()
def __init__(self,
mode,
quality='semantic',
joint_transform_list=None,
img_transform=None,
label_transform=None,
eval_folder=None):
super(Loader, self).__init__(quality=quality,
mode=mode,
joint_transform_list=joint_transform_list,
img_transform=img_transform,
label_transform=label_transform)
root = cfg.DATASET.VAIHINGEN_DIR
self.mode = mode
self.fill_colormap()
######################################################################
# Assemble image lists
######################################################################
if mode == 'folder':
self.all_imgs = make_dataset_folder(eval_folder)
else:
splits = {
'train': 'training',
'val': 'validation',
'trainval': 'trainval',
'test': 'testing'
}
split_name = splits[mode]
#img_ext = 'jpg'
#mask_ext = 'png'
#img_root = os.path.join(root, split_name, 'images')
#mask_root = os.path.join(root, split_name, 'labels')
if mode is 'test':
self.all_imgs = make_dataset(root, quality, mode)
else:
self.all_imgs = make_dataset(root, quality, mode)
logx.msg('all imgs {}'.format(len(self.all_imgs)))
self.centroids = uniform.build_centroids(self.all_imgs,
self.num_classes,
self.train,
cv=cfg.DATASET.CV)
self.build_epoch()
def test_epoch(epoch):
model.eval()
losses = 0.0
total, correct = 0.0, 0.0
with torch.no_grad():
for step, (x, y) in enumerate(val_loader):
x, y = x.to(config.device), y.to(config.device)
out = model(x)
loss = criterion(out, y)
losses += loss.cpu().detach().numpy()
_, pred = torch.max(out.data, 1)
total += y.size(0)
correct += (pred == y).squeeze().sum().cpu().numpy()
save_dict = {
'state_dict': model.state_dict()
}
logx.msg("epoch {} validation loss {} validation acc {}".format(epoch, losses / (step + 1), correct / total))
logx.metric('val', {'loss': losses / (step + 1), 'acc': correct / total})
logx.save_model(save_dict, losses, epoch, higher_better=False, delete_old=True)
def __init__(self,
mode,
quality='semantic',
joint_transform_list=None,
img_transform=None,
label_transform=None,
eval_folder=None):
super(Loader, self).__init__(quality=quality,
mode=mode,
joint_transform_list=joint_transform_list,
img_transform=img_transform,
label_transform=label_transform)
root = cfg.DATASET.MAPILLARY_DIR
config_fn = os.path.join(root, 'config.json')
self.fill_colormap_and_names(config_fn)
######################################################################
# Assemble image lists
######################################################################
if mode == 'folder':
self.all_imgs = make_dataset_folder(eval_folder)
else:
splits = {
'train': 'training',
'val': 'validation',
'test': 'testing'
}
split_name = splits[mode]
img_ext = 'jpg'
mask_ext = 'png'
img_root = os.path.join(root, split_name, 'images')
mask_root = os.path.join(root, split_name, 'labels')
self.all_imgs = self.find_images(img_root, mask_root, img_ext,
mask_ext)
logx.msg('all imgs {}'.format(len(self.all_imgs)))
self.centroids = uniform.build_centroids(self.all_imgs,
self.num_classes,
self.train,
cv=cfg.DATASET.CV)
self.build_epoch()
