def main(config, model_path, cuda, crf, camera_id):
# Configuration
CONFIG = Dict(yaml.load(open(config)))
cuda = cuda and torch.cuda.is_available()
if cuda:
current_device = torch.cuda.current_device()
print('Running on', torch.cuda.get_device_name(current_device))
# Label list
with open(CONFIG.LABELS) as f:
classes = {}
for label in f:
label = label.rstrip().split('\t')
classes[int(label[0])] = label[1].split(',')[0]
# Load a model
state_dict = torch.load(model_path)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
model.load_state_dict(state_dict)
model.eval()
if cuda:
model.cuda()
image_size = (CONFIG.IMAGE.SIZE.TEST, ) * 2
cap = cv2.VideoCapture(camera_id)
cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'YUYV'))
while True:
# Image preprocessing
ret, frame = cap.read()
image = cv2.resize(frame.astype(float), image_size)
raw_image = image.astype(np.uint8)
image -= np.array([
float(CONFIG.IMAGE.MEAN.B),
float(CONFIG.IMAGE.MEAN.G),
float(CONFIG.IMAGE.MEAN.R),
])
image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0)
image = image.cuda() if cuda else image
# Inference
output = model(Variable(image, volatile=True))
output = F.upsample(output, size=image_size, mode='bilinear')
output = F.softmax(output, dim=1)
output = output.data.cpu().numpy()[0]
if crf:
output = dense_crf(raw_image, output)
labelmap = np.argmax(output.transpose(1, 2, 0), axis=2)
labelmap = labelmap.astype(float) / CONFIG.N_CLASSES
labelmap = cm.jet_r(labelmap)[..., :-1] * 255.0
cv2.addWeighted(np.uint8(labelmap), 0.5, raw_image, 0.5, 0.0,
raw_image)
cv2.imshow('DeepLabV2', raw_image)
cv2.waitKey(50)
def main(dataset):
WHITELIST = ['kernel_size', 'stride', 'padding', 'dilation', 'eps', 'momentum'] # NOQA
MODEL_ROOT = 'data/models/deeplab_resnet101/voc12/'
CONFIG = {
'voc12': {
'path_caffe_model': 'data/models/deeplab_resnet101/voc12/train2_iter_20000.caffemodel',
'path_pytorch_model': 'data/models/deeplab_resnet101/voc12/deeplabv2_resnet101_VOC2012.pth',
'n_classes': 21,
},
'coco_init': {
'path_caffe_model': 'data/models/deeplab_resnet101/voc12/init.caffemodel',
'path_pytorch_model': 'data/models/deeplab_resnet101/coco_init/deeplabv2_resnet101_COCO_init.pth',
'n_classes': 91,
},
}.get(dataset)
params = parse_caffemodel(CONFIG['path_caffe_model'])
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG['n_classes'])
model.eval()
own_state = model.state_dict()
state_dict = OrderedDict()
for layer_name, layer_dict in params.items():
for param_name, values in layer_dict.items():
if param_name in WHITELIST:
continue
param_name = translate_layer_name(layer_name) + '.' + param_name
if param_name in own_state:
values = torch.FloatTensor(values)
values = values.view_as(own_state[param_name])
state_dict[param_name] = values
print(layer_name.ljust(25), '->', param_name, ': Copied')
torch.save(state_dict, CONFIG['path_pytorch_model'])
def setup_model(model_path, device, CONFIG):
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
state_dict = torch.load(model_path,
map_location=lambda storage, loc: storage)
model.load_state_dict(state_dict)
model.eval()
model.to(device)
return model
def main(dataset):
WHITELIST = [
'kernel_size', 'stride', 'padding', 'dilation', 'eps', 'momentum'
] # NOQA
CONFIG = {
'voc12': {
'path_caffe_model':
'data/models/deeplab_resnet101/voc12/train2_iter_20000.caffemodel',
'path_pytorch_model':
'data/models/deeplab_resnet101/voc12/deeplabv2_resnet101_VOC2012.pth',
'n_classes': 21,
},
'coco_init': {
'path_caffe_model':
'data/models/deeplab_resnet101/coco_init/init.caffemodel',
'path_pytorch_model':
'data/models/deeplab_resnet101/coco_init/deeplabv2_resnet101_COCO_init.pth',
'n_classes': 91,
},
'init': {
# The same as the coco_init parameters
'path_caffe_model':
'data/models/deeplab_resnet101/init/deeplabv2_resnet101_init.caffemodel',
'path_pytorch_model':
'data/models/deeplab_resnet101/init/deeplabv2_resnet101_init.pth',
'n_classes': 91,
},
}.get(dataset)
params = parse_caffemodel(CONFIG['path_caffe_model'])
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG['n_classes'])
model.eval()
own_state = model.state_dict()
state_dict = OrderedDict()
for layer_name, layer_dict in params.items():
for param_name, values in layer_dict.items():
if param_name in WHITELIST and dataset != 'coco_init' and dataset != 'init':
attribute = translate_layer_name(layer_name)
attribute = eval('model.' + attribute + '.' + param_name)
if isinstance(attribute, tuple):
if attribute[0] != values:
raise ValueError
else:
if abs(attribute - values) > 1e-4:
raise ValueError
print(layer_name.ljust(20), '->', param_name, attribute,
values, ': Checked!')
continue
param_name = translate_layer_name(layer_name) + '.' + param_name
if param_name in own_state:
values = torch.FloatTensor(values)
values = values.view_as(own_state[param_name])
state_dict[param_name] = values
print(layer_name.ljust(20), '->', param_name, ': Copied!')
torch.save(state_dict, CONFIG['path_pytorch_model'])
def setup_model(model_path, CONFIG, train=True):
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)
if train:
model.load_state_dict(state_dict, strict=False) # to skip ASPP
model = nn.DataParallel(model)
else:
model.load_state_dict(state_dict)
model = nn.DataParallel(model)
model.eval()
return model
def loadModel(model_path, config_obj, cuda=True):
cuda = cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
if cuda:
current_device = torch.cuda.current_device()
print("Running on", torch.cuda.get_device_name(current_device))
else:
print("Running on CPU")
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=config_obj.N_CLASSES)
state_dict = torch.load(model_path,
map_location=lambda storage, loc: storage)
model.load_state_dict(state_dict)
model.eval()
model.to(device)
return model
def __init__(self,
device,
resnet,
pseudolabeling,
spnetcheckpoint,
pretrained,
model_path,
CONFIG,
datamanager,
lr,
bn="BN",
continue_from=None,
mirroring=False,
scale=-1):
_, _, seen_novel_classes, _, visible_classes, _, _ = datamanager.get_Classes(
)
class_emb, to_ignore_class_emb, _ = datamanager.get_clsEmbs()
#MODELS
#pseudo label generator
if scale == 0:
scales = 0
elif scale == 1:
scales = [0.5, 0.75]
elif scale == 2:
scales = [1.5, 1.75]
else:
scales = None
if scales is None and mirroring == False:
model_old = DeepLabV2_ResNet101_MSC(class_emb.shape[1],
class_emb[seen_novel_classes],
preload='train',
resnet=resnet)
else:
model_old = DeepLabV2_ResNet101_consistency(
class_emb.shape[1],
class_emb[seen_novel_classes],
preload='train',
resnet=resnet,
scales=scales,
mirroring=mirroring)
# under training model
if pseudolabeling >= 0:
model = DeepLabV2_ResNet101_MSC(class_emb.shape[1],
class_emb[seen_novel_classes],
preload='train',
resnet=resnet)
else:
model = DeepLabV2_ResNet101_MSC(class_emb.shape[1],
class_emb[visible_classes],
preload='train',
resnet=resnet)
if not spnetcheckpoint:
if pretrained == True:
print("Loading pretrained")
state_dict = torch.load(CONFIG.INIT_MODEL_O,
map_location='cpu')
model_old.load_state_dict(state_dict, strict=True)
del state_dict
model_old.base.aspp = _ASPP(2048, class_emb.shape[1], [6, 12, 18, 24])
model.base.aspp = _ASPP(2048, class_emb.shape[1], [6, 12, 18, 24])
if spnetcheckpoint == True:
if model_path is None:
model_path = CONFIG.INIT_INITIAL_CHECKPOINT
p = torch.load(model_path, map_location='cpu')
new_state_dict = OrderedDict()
for k, v in p['state_dict'].items():
name = k[7:] # remove `module.`
if resnet == 'spnet':
name = name.replace("scale", "base") # 'scale'->base
name = name.replace("stages.", "")
new_state_dict[name] = v
model_old.load_state_dict(new_state_dict, strict=True)
if continue_from is None:
model.base = copy.deepcopy(model_old.base)
if pseudolabeling >= 0:
model.base.setClassEmb(class_emb[seen_novel_classes])
else:
model.base.setClassEmb(class_emb[visible_classes])
model_old.base.setClassEmb(to_ignore_class_emb)
self.model_old = DistributedDataParallel(
model_old.to(device), device_ids=[torch.distributed.get_rank()])
self.model = DistributedDataParallel(
model.to(device), device_ids=[torch.distributed.get_rank()])
if continue_from is not None:
state_dict = torch.load(continue_from, map_location='cpu')
self.model.load_state_dict(state_dict['state_dict'], strict=True)
# Optimizer
self.optimizer = {
"sgd":
torch.optim.SGD(
# cf lr_mult and decay_mult in train.prototxt
params=[{
"params": get_params(self.model, key="1x"),
"lr": lr,
"weight_decay": CONFIG.WEIGHT_DECAY,
}, {
"params": get_params(self.model, key="10x"),
"lr": 10 * lr,
"weight_decay": CONFIG.WEIGHT_DECAY,
}, {
"params": get_params(self.model, key="20x"),
"lr": 20 * lr,
"weight_decay": 0.0,
}],
momentum=CONFIG.MOMENTUM,
),
"adam":
torch.optim.Adam(
# cf lr_mult and decay_mult in train.prototxt
params=[{
"params": get_params(self.model, key="1x"),
"lr": lr,
"weight_decay": CONFIG.WEIGHT_DECAY,
}, {
"params": get_params(self.model, key="10x"),
"lr": 10 * lr,
"weight_decay": CONFIG.WEIGHT_DECAY,
}, {
"params": get_params(self.model, key="20x"),
"lr": 20 * lr,
"weight_decay": 0.0,
}])
# Add any other optimizer
}.get(CONFIG.OPTIMIZER)
if continue_from is not None:
if 'optimizer' in state_dict:
self.optimizer.load_state_dict(state_dict['optimizer'])
else:
print("[continue from] optimizer not defined")
del state_dict
def main(config, cuda, gpu):
# Configuration
CONFIG = Dict(yaml.load(open(config)))
# CUDA check
cuda = cuda and torch.cuda.is_available()
if cuda:
gpu_ids = [int(string) for string in gpu.split(',')]
current_device = torch.cuda.current_device()
print('Running on', torch.cuda.get_device_name(current_device),
gpu_ids)
# Dataset
dataset = CocoStuff10k(
root=CONFIG.ROOT,
split='train',
image_size=513,
crop_size=CONFIG.IMAGE.SIZE.TRAIN,
scale=True,
flip=True,
)
# DataLoader
loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=CONFIG.BATCH_SIZE,
num_workers=CONFIG.NUM_WORKERS,
shuffle=True,
)
loader_iter = iter(loader)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
state_dict = torch.load(CONFIG.INIT_MODEL)
model.load_state_dict(state_dict, strict=False) # Skip "aspp" layer
model = nn.DataParallel(model, device_ids=gpu_ids)
if cuda:
model.cuda()
# Optimizer
optimizer = {
'sgd':
torch.optim.SGD(
# cf lr_mult and decay_mult in train.prototxt
params=[{
'params': get_lr_params(model.module, key='1x'),
'lr': CONFIG.LR,
'weight_decay': CONFIG.WEIGHT_DECAY
}, {
'params': get_lr_params(model.module, key='10x'),
'lr': 10 * CONFIG.LR,
'weight_decay': CONFIG.WEIGHT_DECAY
}, {
'params': get_lr_params(model.module, key='20x'),
'lr': 20 * CONFIG.LR,
'weight_decay': 0.0
}],
momentum=CONFIG.MOMENTUM,
),
}.get(CONFIG.OPTIMIZER)
# Loss definition
criterion = CrossEntropyLoss2d(ignore_index=CONFIG.IGNORE_LABEL)
if cuda:
criterion.cuda()
# TensorBoard Logger
writer = SummaryWriter(CONFIG.LOG_DIR)
loss_meter = MovingAverageValueMeter(20)
model.train()
model.module.scale.freeze_bn()
for iteration in tqdm(
range(1, CONFIG.ITER_MAX + 1),
total=CONFIG.ITER_MAX,
leave=False,
dynamic_ncols=True,
):
# Set a learning rate
poly_lr_scheduler(
optimizer=optimizer,
init_lr=CONFIG.LR,
iter=iteration - 1,
lr_decay_iter=CONFIG.LR_DECAY,
max_iter=CONFIG.ITER_MAX,
power=CONFIG.POLY_POWER,
)
# Clear gradients (ready to accumulate)
optimizer.zero_grad()
iter_loss = 0
for i in range(1, CONFIG.ITER_SIZE + 1):
data, target = next(loader_iter)
# Image
data = data.cuda() if cuda else data
data = Variable(data)
# Propagate forward
outputs = model(data)
# Loss
loss = 0
for output in outputs:
# Resize target for {100%, 75%, 50%, Max} outputs
target_ = resize_target(target, output.size(2))
target_ = target_.cuda() if cuda else target_
target_ = Variable(target_)
# Compute crossentropy loss
loss += criterion(output, target_)
# Backpropagate (just compute gradients wrt the loss)
loss /= float(CONFIG.ITER_SIZE)
loss.backward()
iter_loss += loss.data[0]
# Reload dataloader
if ((iteration - 1) * CONFIG.ITER_SIZE + i) % len(loader) == 0:
loader_iter = iter(loader)
loss_meter.add(iter_loss)
# Update weights with accumulated gradients
optimizer.step()
# TensorBoard
if iteration % CONFIG.ITER_TF == 0:
writer.add_scalar('train_loss', loss_meter.value()[0], iteration)
for i, o in enumerate(optimizer.param_groups):
writer.add_scalar('train_lr_group{}'.format(i), o['lr'],
iteration)
if iteration % 1000 != 0:
continue
for name, param in model.named_parameters():
name = name.replace('.', '/')
writer.add_histogram(name, param, iteration, bins="auto")
if param.requires_grad:
writer.add_histogram(name + '/grad',
param.grad,
iteration,
bins="auto")
# Save a model
if iteration % CONFIG.ITER_SNAP == 0:
torch.save(
model.module.state_dict(),
osp.join(CONFIG.SAVE_DIR,
'checkpoint_{}.pth'.format(iteration)),
)
# Save a model
if iteration % 100 == 0:
torch.save(
model.module.state_dict(),
osp.join(CONFIG.SAVE_DIR, 'checkpoint_current.pth'),
)
torch.save(
model.module.state_dict(),
osp.join(CONFIG.SAVE_DIR, 'checkpoint_final.pth'),
)
def main(config, excludeval, embedding, model_path, run, cuda, crf, redo,
imagedataset, threshold):
pth_extn = '.pth.tar'
if osp.isfile(model_path.replace(
pth_extn, "_" + run + ".json")) and not threshold and not redo:
print("Already Done!")
with open(model_path.replace(pth_extn,
"_" + run + ".json")) as json_file:
data = json.load(json_file)
for key, value in data.items():
if not key == "Class IoU":
print(key, value)
sys.exit()
cuda = cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
if cuda:
current_device = torch.cuda.current_device()
print("Running on", torch.cuda.get_device_name(current_device))
else:
print("Running on CPU")
# Configuration
CONFIG = Dict(yaml.load(open(config)))
datadir = os.path.join('data/datasets', imagedataset)
print("Split dir: ", datadir)
savedir = osp.dirname(model_path)
epoch = re.findall("checkpoint_(.*)\." + pth_extn[1:],
osp.basename(model_path))[-1]
val = None
visible_classes = None
if run == 'zlss' or run == 'flss':
val = np.load(datadir + '/split/test_list.npy')
visible_classes = np.load(datadir + '/split/novel_cls.npy')
elif run == 'gzlss' or run == 'gflss':
val = np.load(datadir + '/split/test_list.npy')
if excludeval:
vals_cls = np.asarray(np.load(datadir + '/split/seen_cls.npy'),
dtype=int)
else:
vals_cls = np.asarray(np.concatenate([
np.load(datadir + '/split/seen_cls.npy'),
np.load(datadir + '/split/val_cls.npy')
]),
dtype=int)
valu_cls = np.load(datadir + '/split/novel_cls.npy')
visible_classes = np.concatenate([vals_cls, valu_cls])
else:
print("invalid run ", run)
sys.exit()
if threshold is not None and run != 'gzlss':
print("invalid run for threshold", run)
sys.exit()
cls_map = np.array([255] * 256)
for i, n in enumerate(visible_classes):
cls_map[n] = i
if threshold is not None:
savedir = osp.join(savedir, str(threshold))
if crf is not None:
savedir = savedir + '-crf'
if run == 'gzlss' or run == 'gflss':
novel_cls_map = np.array([255] * 256)
for i, n in enumerate(list(valu_cls)):
novel_cls_map[cls_map[n]] = i
seen_cls_map = np.array([255] * 256)
for i, n in enumerate(list(vals_cls)):
seen_cls_map[cls_map[n]] = i
if threshold is not None:
thresholdv = np.asarray(np.zeros((visible_classes.shape[0], 1)),
dtype=np.float)
thresholdv[np.in1d(visible_classes, vals_cls), 0] = threshold
thresholdv = torch.tensor(thresholdv).float().cuda()
visible_classesp = np.concatenate([visible_classes, [255]])
all_labels = np.genfromtxt(datadir + '/labels_2.txt',
delimiter='\t',
usecols=1,
dtype='str')
print("Visible Classes: ", visible_classes)
# Dataset
dataset = get_dataset(CONFIG.DATASET)(
train=None,
test=val,
root=CONFIG.ROOT,
split=CONFIG.SPLIT.TEST,
base_size=CONFIG.IMAGE.SIZE.TEST,
mean=(CONFIG.IMAGE.MEAN.B, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.R),
warp=CONFIG.WARP_IMAGE,
scale=None,
flip=False,
)
if embedding == 'word2vec':
class_emb = pickle.load(
open(datadir + '/word_vectors/word2vec.pkl', "rb"))
elif embedding == 'fasttext':
class_emb = pickle.load(
open(datadir + '/word_vectors/fasttext.pkl', "rb"))
elif embedding == 'fastnvec':
class_emb = np.concatenate([
pickle.load(open(datadir + '/word_vectors/fasttext.pkl', "rb")),
pickle.load(open(datadir + '/word_vectors/word2vec.pkl', "rb"))
],
axis=1)
else:
print("invalid emb ", embedding)
sys.exit()
class_emb = class_emb[visible_classes]
class_emb = F.normalize(torch.tensor(class_emb), p=2, dim=1).cuda()
print("Embedding dim: ", class_emb.shape[1])
print("# Visible Classes: ", class_emb.shape[0])
# DataLoader
loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=CONFIG.BATCH_SIZE.TEST,
num_workers=CONFIG.NUM_WORKERS,
shuffle=False,
)
torch.set_grad_enabled(False)
# Model
model = DeepLabV2_ResNet101_MSC(class_emb.shape[1], class_emb)
sdir = osp.join(savedir, model_path.replace(pth_extn, ""), str(epoch), run)
state_dict = torch.load(model_path,
map_location=lambda storage, loc: storage)
model = nn.DataParallel(model)
model.load_state_dict(state_dict['state_dict'])
model.eval()
model.to(device)
imgfeat = []
targets, outputs = [], []
for data, target, img_id in tqdm(loader,
total=len(loader),
leave=False,
dynamic_ncols=True):
# Image
data = data.to(device)
# Forward propagation
output = model(data)
output = F.interpolate(output,
size=data.shape[2:],
mode="bilinear",
align_corners=False)
output = F.softmax(output, dim=1)
if threshold is not None:
output = output - thresholdv.view(1, -1, 1, 1)
target = cls_map[target.numpy()]
# Postprocessing
if crf:
output = output.data.cpu().numpy()
crf_output = np.zeros(output.shape)
images = data.data.cpu().numpy().astype(np.uint8)
for i, (image, prob_map) in enumerate(zip(images, output)):
image = image.transpose(1, 2, 0)
crf_output[i] = dense_crf(image, prob_map)
output = crf_output
output = np.argmax(output, axis=1)
else:
output = torch.argmax(output, dim=1).cpu().numpy()
for o, t in zip(output, target):
outputs.append(o)
targets.append(t)
if run == 'gzlss' or run == 'gflss':
score, class_iou = scores_gzsl(targets,
outputs,
n_class=len(visible_classes),
seen_cls=cls_map[vals_cls],
unseen_cls=cls_map[valu_cls])
else:
score, class_iou = scores(targets,
outputs,
n_class=len(visible_classes))
for k, v in score.items():
print(k, v)
score["Class IoU"] = {}
for i in range(len(visible_classes)):
score["Class IoU"][all_labels[visible_classes[i]]] = class_iou[i]
if threshold is not None:
with open(
model_path.replace(pth_extn, "_" + run + '_T' +
str(threshold) + ".json"), "w") as f:
json.dump(score, f, indent=4, sort_keys=True)
else:
with open(model_path.replace(pth_extn, "_" + run + ".json"), "w") as f:
json.dump(score, f, indent=4, sort_keys=True)
print(score["Class IoU"])
def main(dataset):
WHITELIST = [
"kernel_size", "stride", "padding", "dilation", "eps", "momentum"
]
CONFIG = {
"voc12": {
"path_caffe_model":
"data/models/deeplab_resnet101/voc12/train2_iter_20000.caffemodel",
"path_pytorch_model":
"data/models/deeplab_resnet101/voc12/deeplabv2_resnet101_VOC2012.pth",
"n_classes": 21,
},
"coco_init": {
"path_caffe_model":
"data/models/deeplab_resnet101/coco_init/init.caffemodel",
"path_pytorch_model":
"data/models/deeplab_resnet101/coco_init/deeplabv2_resnet101_COCO_init.pth",
"n_classes": 91,
},
"init": {
# The same as the coco_init parameters
"path_caffe_model":
"data/models/deeplab_resnet101/init/deeplabv2_resnet101_init.caffemodel",
"path_pytorch_model":
"data/models/deeplab_resnet101/init/deeplabv2_resnet101_init.pth",
"n_classes": 91,
},
}.get(dataset)
params = parse_caffemodel(CONFIG["path_caffe_model"])
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG["n_classes"])
model.eval()
own_state = model.state_dict()
state_dict = OrderedDict()
for layer_name, layer_dict in params.items():
for param_name, values in layer_dict.items():
if param_name in WHITELIST and dataset != "coco_init" and dataset != "init":
attribute = translate_layer_name(layer_name)
attribute = eval("model." + attribute + "." + param_name)
if isinstance(attribute, tuple):
if attribute[0] != values:
raise ValueError
else:
if abs(attribute - values) > 1e-4:
raise ValueError
print(
layer_name.ljust(20),
"->",
param_name,
attribute,
values,
": Checked!",
)
continue
param_name = translate_layer_name(layer_name) + "." + param_name
if param_name in own_state:
values = torch.FloatTensor(values)
values = values.view_as(own_state[param_name])
state_dict[param_name] = values
print(layer_name.ljust(20), "->", param_name, ": Copied!")
torch.save(state_dict, CONFIG["path_pytorch_model"])
def main(dataset):
WHITELIST = ["kernel_size", "stride", "padding", "dilation", "eps", "momentum"]
CONFIG = {
"voc12": {
"path_caffe_model": "data/models/deeplab_resnet101/voc12/train2_iter_20000.caffemodel",
"path_pytorch_model": "data/models/deeplab_resnet101/voc12/deeplabv2_resnet101_VOC2012.pth",
"n_classes": 21,
},
"init": {
"path_caffe_model": "data/models/deeplab_resnet101/coco_init/init.caffemodel",
"path_pytorch_model": "data/models/deeplab_resnet101/coco_init/deeplabv2_resnet101_COCO_init.pth",
"n_classes": 91,
},
}.get(dataset)
params = parse_caffemodel(CONFIG["path_caffe_model"])
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG["n_classes"])
model.eval()
own_state = model.state_dict()
rel_tol = 1e-7
state_dict = OrderedDict()
for layer_name, layer_dict in params.items():
for param_name, values in layer_dict.items():
if param_name in WHITELIST and dataset != "coco_init" and dataset != "init":
attribute = translate_layer_name(layer_name)
attribute = eval("model." + attribute + "." + param_name)
if isinstance(attribute, tuple):
assert (
attribute[0] == values
), "Inconsistent values: {}@{}, {}@{}".format(
attribute[0],
translate_layer_name(layer_name) + "." + param_name,
values,
layer_name,
)
else:
assert (
abs(attribute - values) < rel_tol
), "Inconsistent values: {}@{}, {}@{}".format(
attribute,
translate_layer_name(layer_name) + "." + param_name,
values,
layer_name,
)
print(
layer_name.ljust(20),
"->",
param_name,
attribute,
values,
": Checked!",
)
continue
param_name = translate_layer_name(layer_name) + "." + param_name
if param_name in own_state:
values = torch.FloatTensor(values)
values = values.view_as(own_state[param_name])
state_dict[param_name] = values
print(layer_name.ljust(20), "->", param_name, ": Copied!")
try:
print("\033[32mVerify the converted model\033[00m")
model.load_state_dict(state_dict)
except:
import traceback
traceback.print_exc()
print("\033[32mVerify with ignoring ASPP (strict=False)\033[00m")
model.load_state_dict(state_dict, strict=False)
print("Saving to", CONFIG["path_pytorch_model"])
torch.save(state_dict, CONFIG["path_pytorch_model"])
def main(config, cuda):
device = torch.device("cuda" if cuda and torch.cuda.is_available() else "cpu")
if cuda:
current_device = torch.cuda.current_device()
print("Running on", torch.cuda.get_device_name(current_device))
else:
print("Running on CPU")
# Configuration
CONFIG = Dict(yaml.load(open(config)))
# Dataset
dataset = CocoStuff10k(
root=CONFIG.ROOT,
split="train",
image_size=513,
crop_size=CONFIG.IMAGE.SIZE.TRAIN,
scale=True,
flip=True,
)
# DataLoader
loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=CONFIG.BATCH_SIZE,
num_workers=CONFIG.NUM_WORKERS,
shuffle=True,
)
loader_iter = iter(loader)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
state_dict = torch.load(CONFIG.INIT_MODEL)
model.load_state_dict(state_dict, strict=False) # Skip "aspp" layer
model = nn.DataParallel(model)
model.to(device)
# Optimizer
optimizer = {
"sgd": torch.optim.SGD(
# cf lr_mult and decay_mult in train.prototxt
params=[
{
"params": get_lr_params(model.module, key="1x"),
"lr": CONFIG.LR,
"weight_decay": CONFIG.WEIGHT_DECAY,
},
{
"params": get_lr_params(model.module, key="10x"),
"lr": 10 * CONFIG.LR,
"weight_decay": CONFIG.WEIGHT_DECAY,
},
{
"params": get_lr_params(model.module, key="20x"),
"lr": 20 * CONFIG.LR,
"weight_decay": 0.0,
},
],
momentum=CONFIG.MOMENTUM,
)
}.get(CONFIG.OPTIMIZER)
# Loss definition
criterion = CrossEntropyLoss2d(ignore_index=CONFIG.IGNORE_LABEL)
criterion.to(device)
# TensorBoard Logger
writer = SummaryWriter(CONFIG.LOG_DIR)
loss_meter = MovingAverageValueMeter(20)
model.train()
model.module.scale.freeze_bn()
for iteration in tqdm(
range(1, CONFIG.ITER_MAX + 1),
total=CONFIG.ITER_MAX,
leave=False,
dynamic_ncols=True,
):
# Set a learning rate
poly_lr_scheduler(
optimizer=optimizer,
init_lr=CONFIG.LR,
iter=iteration - 1,
lr_decay_iter=CONFIG.LR_DECAY,
max_iter=CONFIG.ITER_MAX,
power=CONFIG.POLY_POWER,
)
# Clear gradients (ready to accumulate)
optimizer.zero_grad()
iter_loss = 0
for i in range(1, CONFIG.ITER_SIZE + 1):
try:
data, target = next(loader_iter)
except:
loader_iter = iter(loader)
data, target = next(loader_iter)
# Image
data = data.to(device)
# Propagate forward
outputs = model(data)
# Loss
loss = 0
for output in outputs:
# Resize target for {100%, 75%, 50%, Max} outputs
target_ = resize_target(target, output.size(2))
target_ = target_.to(device)
# Compute crossentropy loss
loss += criterion(output, target_)
# Backpropagate (just compute gradients wrt the loss)
loss /= float(CONFIG.ITER_SIZE)
loss.backward()
iter_loss += float(loss)
loss_meter.add(iter_loss)
# Update weights with accumulated gradients
optimizer.step()
# TensorBoard
if iteration % CONFIG.ITER_TF == 0:
writer.add_scalar("train_loss", loss_meter.value()[0], iteration)
for i, o in enumerate(optimizer.param_groups):
writer.add_scalar("train_lr_group{}".format(i), o["lr"], iteration)
# for name, param in model.named_parameters():
# name = name.replace('.', '/')
# writer.add_histogram(name, param, iteration, bins="auto")
# if param.requires_grad:
# writer.add_histogram(name + '/grad', param.grad, iteration, bins="auto")
# Save a model
if iteration % CONFIG.ITER_SNAP == 0:
torch.save(
model.module.state_dict(),
osp.join(CONFIG.SAVE_DIR, "checkpoint_{}.pth".format(iteration)),
)
# Save a model
if iteration % 100 == 0:
torch.save(
model.module.state_dict(),
osp.join(CONFIG.SAVE_DIR, "checkpoint_current.pth"),
)
torch.save(
model.module.state_dict(), osp.join(CONFIG.SAVE_DIR, "checkpoint_final.pth")
)
def main(config, model_path, cuda, crf, camera_id):
cuda = cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
if cuda:
current_device = torch.cuda.current_device()
print("Running on", torch.cuda.get_device_name(current_device))
else:
print("Running on CPU")
# Configuration
CONFIG = Dict(yaml.load(open(config)))
# Label list
with open(CONFIG.LABELS) as f:
classes = {}
for label in f:
label = label.rstrip().split("\t")
classes[int(label[0])] = label[1].split(",")[0]
torch.set_grad_enabled(False)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
model.load_state_dict(torch.load(model_path))
model.eval()
model.to(device)
cap = cv2.VideoCapture(camera_id)
cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*"YUYV"))
while True:
# Image preprocessing
ret, frame = cap.read()
h, w, c = frame.shape
image = frame.astype(np.float32)
image -= np.array([
float(CONFIG.IMAGE.MEAN.B),
float(CONFIG.IMAGE.MEAN.G),
float(CONFIG.IMAGE.MEAN.R),
])
image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0)
image = image.to(device)
# Inference
output = model(image)
output = F.upsample(output,
size=(h, w),
mode="bilinear",
align_corners=False)
output = F.softmax(output, dim=1)
output = output.data.cpu().numpy()[0]
if crf:
output = dense_crf(raw_image, output)
labelmap = np.argmax(output.transpose(1, 2, 0), axis=2)
labelmap = labelmap.astype(float) / CONFIG.N_CLASSES
labelmap = cm.jet_r(labelmap)[..., :-1] * 255.0
cv2.addWeighted(np.uint8(labelmap), 0.5, frame, 0.5, 0.0, frame)
cv2.imshow("DeepLabV2", frame)
cv2.waitKey(10)
def main(config, model_path, cuda, crf):
cuda = cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
if cuda:
current_device = torch.cuda.current_device()
print("Running on", torch.cuda.get_device_name(current_device))
else:
print("Running on CPU")
# Configuration
CONFIG = Dict(yaml.load(open(config)))
# Dataset 10k or 164k
dataset = get_dataset(CONFIG.DATASET)(
root=CONFIG.ROOT,
split=CONFIG.SPLIT.VAL,
base_size=CONFIG.IMAGE.SIZE.TEST,
mean=(CONFIG.IMAGE.MEAN.B, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.R),
warp=CONFIG.WARP_IMAGE,
scale=None,
flip=False,
)
# DataLoader
loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=CONFIG.BATCH_SIZE.TEST,
num_workers=CONFIG.NUM_WORKERS,
shuffle=False,
)
torch.set_grad_enabled(False)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
state_dict = torch.load(model_path,
map_location=lambda storage, loc: storage)
model.load_state_dict(state_dict)
model = nn.DataParallel(model)
model.eval()
model.to(device)
targets, outputs = [], []
for data, target in tqdm(loader,
total=len(loader),
leave=False,
dynamic_ncols=True):
# Image
data = data.to(device)
# Forward propagation
output = model(data)
output = F.interpolate(output, size=data.shape[2:], mode="bilinear")
output = F.softmax(output, dim=1)
output = output.data.cpu().numpy()
# Postprocessing
if crf:
crf_output = np.zeros(output.shape)
images = data.data.cpu().numpy().astype(np.uint8)
for i, (image, prob_map) in enumerate(zip(images, output)):
image = image.transpose(1, 2, 0)
crf_output[i] = dense_crf(image, prob_map)
output = crf_output
output = np.argmax(output, axis=1)
target = target.numpy()
for o, t in zip(output, target):
outputs.append(o)
targets.append(t)
score, class_iou = scores(targets, outputs, n_class=CONFIG.N_CLASSES)
for k, v in score.items():
print(k, v)
score["Class IoU"] = {}
for i in range(CONFIG.N_CLASSES):
score["Class IoU"][i] = class_iou[i]
with open(model_path.replace(".pth", ".json"), "w") as f:
json.dump(score, f, indent=4, sort_keys=True)
def main(config, cuda):
# Configuration
with open(config) as f:
CONFIG = yaml.load(f)
cuda = cuda and torch.cuda.is_available()
# Dataset
dataset = get_dataset(CONFIG['DATASET'])(
root=CONFIG['ROOT'],
split='train',
image_size=(CONFIG['IMAGE']['SIZE']['TRAIN'],
CONFIG['IMAGE']['SIZE']['TRAIN']),
scale=True,
flip=True,
# preload=True
)
# DataLoader
loader = torch.utils.data.DataLoader(dataset=dataset,
batch_size=CONFIG['BATCH_SIZE'],
num_workers=CONFIG['NUM_WORKERS'],
shuffle=True)
loader_iter = iter(loader)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG['N_CLASSES'])
state_dict = torch.load(CONFIG['INIT_MODEL'])
model.load_state_dict(state_dict, strict=False) # Skip "aspp" layer
if cuda:
model.cuda()
# Optimizer
optimizer = {
'sgd':
torch.optim.SGD(
params=[
{
'params': get_1x_lr_params(model),
'lr': float(CONFIG['LR'])
},
{
'params': get_10x_lr_params(model),
'lr': 10 * float(CONFIG['LR'])
} # NOQA
],
lr=float(CONFIG['LR']),
momentum=float(CONFIG['MOMENTUM']),
weight_decay=float(CONFIG['WEIGHT_DECAY'])),
}.get(CONFIG['OPTIMIZER'])
# Loss definition
criterion = CrossEntropyLoss2d(ignore_index=CONFIG['IGNORE_LABEL'])
if cuda:
criterion.cuda()
# TensorBoard Logger
writer = SummaryWriter(CONFIG['LOG_DIR'])
loss_meter = MovingAverageValueMeter(20)
model.train()
for iteration in tqdm(range(1, CONFIG['ITER_MAX'] + 1),
total=CONFIG['ITER_MAX'],
leave=False,
dynamic_ncols=True):
# Polynomial lr decay
poly_lr_scheduler(optimizer=optimizer,
init_lr=float(CONFIG['LR']),
iter=iteration - 1,
lr_decay_iter=CONFIG['LR_DECAY'],
max_iter=CONFIG['ITER_MAX'],
power=CONFIG['POLY_POWER'])
optimizer.zero_grad()
iter_loss = 0
for i in range(1, CONFIG['ITER_SIZE'] + 1):
data, target = next(loader_iter)
# Image
data = data.cuda() if cuda else data
data = Variable(data)
# Forward propagation
outputs = model(data)
# Label
target = resize_target(target, outputs[0].size(2))
target = target.cuda() if cuda else target
target = Variable(target)
# Aggregate losses for [100%, 75%, 50%, Max]
loss = 0
for output in outputs:
loss += criterion(output, target)
loss /= CONFIG['ITER_SIZE']
iter_loss += loss.data[0]
loss.backward()
# Reload dataloader
if ((iteration - 1) * CONFIG['ITER_SIZE'] + i) % len(loader) == 0:
loader_iter = iter(loader)
loss_meter.add(iter_loss)
# Back propagation
optimizer.step()
# TensorBoard
if iteration % CONFIG['ITER_TF'] == 0:
writer.add_scalar('train_loss', loss_meter.value()[0], iteration)
# Save a model
if iteration % CONFIG['ITER_SNAP'] == 0:
torch.save(
model.state_dict(),
osp.join(CONFIG['SAVE_DIR'],
'checkpoint_{}.pth.tar'.format(iteration))) # NOQA
writer.add_text('log', 'Saved a model', iteration)
torch.save(model.state_dict(),
osp.join(CONFIG['SAVE_DIR'], 'checkpoint_final.pth.tar'))
size_to_str(u.size()),
fillcolor='lightblue')
else:
dot.node(str(id(var)),
str(type(var).__name__.replace('Backward', '')))
seen.add(var)
if hasattr(var, 'next_functions'):
for u in var.next_functions:
if u[0] is not None:
dot.edge(str(id(u[0])), str(id(var)))
add_nodes(u[0])
if hasattr(var, 'saved_tensors'):
for t in var.saved_tensors:
dot.edge(str(id(t)), str(id(var)))
add_nodes(t)
add_nodes(var.grad_fn)
return dot
model = DeepLabV2_ResNet101_MSC(n_classes=183)
input = Variable(torch.randn(1, 3, 513, 513))
# y = model(input)
# g = make_dot(y, model.state_dict())
# g.view(filename='model', cleanup=True)
with SummaryWriter('runs/graph', comment='DeepLabV2') as w:
w.add_graph(model, (input, ))
def main(dataset, image_path, model_path, cuda, crf):
CONFIG = {
'voc12': {
'path_pytorch_model':
'data/models/deeplab_resnet101/voc12/deeplabv2_resnet101_VOC2012.pth',
'label_list': 'data/datasets/voc12/labels.txt',
'n_classes': 21,
'n_blocks': [3, 4, 23, 3],
'pyramids': [6, 3, 2, 1],
'image': {
'size': {
'train': 473,
'test': 473,
},
'mean': {
'R': 122.675,
'G': 116.669,
'B': 104.008,
}
},
},
'cocostuff': {
'path_pytorch_model':
'data/models/deeplab_resnet101/cocostuff_rgb/deeplabv2_resnet101_VOC2012.pth',
'label_list': 'data/datasets/voc12/labels.txt',
'n_classes': 183,
'n_blocks': [3, 4, 23, 3],
'pyramids': [6, 3, 2, 1],
'image': {
'size': {
'train': 321,
'test': 513,
},
'mean': {
'R': 122.675,
'G': 116.669,
'B': 104.008,
}
},
},
}.get(dataset)
cuda = cuda and torch.cuda.is_available()
# Label list
with open(CONFIG['label_list']) as f:
classes = {}
for label in f:
label = label.rstrip().split('\t')
classes[int(label[0])] = label[1].split(',')[0]
# Load a model
if model_path is None:
state_dict = torch.load(CONFIG['path_pytorch_model'])
else:
state_dict = torch.load(model_path)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG['n_classes'])
model.load_state_dict(state_dict)
model.eval()
if cuda:
model.cuda()
image_size = (CONFIG['image']['size']['test'],
CONFIG['image']['size']['test'])
# Image preprocessing
image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float)
image = cv2.resize(image, image_size)
image_original = image.astype(np.uint8)
image -= np.array([
float(CONFIG['image']['mean']['B']),
float(CONFIG['image']['mean']['G']),
float(CONFIG['image']['mean']['R'])
])
image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0)
image = image.cuda() if cuda else image
# Inference
output = model(Variable(image, volatile=True))
output = F.upsample(output, size=image_size, mode='bilinear')
output = F.softmax(output, dim=1)
output = output.data.cpu().numpy()[0]
if crf:
output = dense_crf(image_original, output)
labelmap = np.argmax(output.transpose(1, 2, 0), axis=2)
labels = np.unique(labelmap)
# Show results
rows = np.floor(np.sqrt(len(labels) + 1))
cols = np.ceil((len(labels) + 1) / rows)
plt.figure(figsize=(10, 10))
ax = plt.subplot(rows, cols, 1)
ax.set_title('Input image')
ax.imshow(image_original[:, :, ::-1])
ax.set_xticks([])
ax.set_yticks([])
for i, label in enumerate(labels):
print '{0:3d}: {1}'.format(label, classes[label])
mask = labelmap == label
ax = plt.subplot(rows, cols, i + 2)
ax.set_title(classes[label])
ax.imshow(np.dstack((mask, ) * 3) * image_original[:, :, ::-1])
ax.set_xticks([])
ax.set_yticks([])
plt.show()
def main(config, image_path, model_path, cuda, crf):
cuda = cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
if cuda:
current_device = torch.cuda.current_device()
print("Running on", torch.cuda.get_device_name(current_device))
else:
print("Running on CPU")
# Configuration
CONFIG = Dict(yaml.load(open(config)))
# Label list
with open(CONFIG.LABELS) as f:
classes = {}
for label in f:
label = label.rstrip().split("\t")
classes[int(label[0])] = label[1].split(",")[0]
torch.set_grad_enabled(False)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
state_dict = torch.load(model_path,
map_location=lambda storage, loc: storage)
model.load_state_dict(state_dict)
model.eval()
model.to(device)
# Image preprocessing
image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float)
scale = CONFIG.IMAGE.SIZE.TEST / max(image.shape[:2])
image = cv2.resize(image, dsize=None, fx=scale, fy=scale)
image_original = image.astype(np.uint8)
image -= np.array([
float(CONFIG.IMAGE.MEAN.B),
float(CONFIG.IMAGE.MEAN.G),
float(CONFIG.IMAGE.MEAN.R),
])
image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0)
image = image.to(device)
# Inference
output = model(image)
output = F.interpolate(output,
size=image.shape[2:],
mode="bilinear",
align_corners=True)
output = F.softmax(output, dim=1)
output = output.data.cpu().numpy()[0]
if crf:
output = dense_crf(image_original, output)
labelmap = np.argmax(output, axis=0)
labels = np.unique(labelmap)
# Show results
rows = np.floor(np.sqrt(len(labels) + 1))
cols = np.ceil((len(labels) + 1) / rows)
plt.figure(figsize=(10, 10))
ax = plt.subplot(rows, cols, 1)
ax.set_title("Input image")
ax.imshow(image_original[:, :, ::-1])
ax.set_xticks([])
ax.set_yticks([])
for i, label in enumerate(labels):
print("{0:3d}: {1}".format(label, classes[label]))
mask = labelmap == label
ax = plt.subplot(rows, cols, i + 2)
ax.set_title(classes[label])
ax.imshow(image_original[..., ::-1])
ax.imshow(mask.astype(np.float32), alpha=0.5, cmap="viridis")
ax.set_xticks([])
ax.set_yticks([])
plt.tight_layout()
plt.show()
def main(config, cuda):
cuda = cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
if cuda:
current_device = torch.cuda.current_device()
print("Running on", torch.cuda.get_device_name(current_device))
else:
print("Running on CPU")
# Configuration
CONFIG = Dict(yaml.load(open(config)))
# Dataset 10k or 164k
dataset = get_dataset(CONFIG.DATASET)(
root=CONFIG.ROOT,
split=CONFIG.SPLIT.TRAIN,
base_size=513,
crop_size=CONFIG.IMAGE.SIZE.TRAIN,
mean=(CONFIG.IMAGE.MEAN.B, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.R),
warp=CONFIG.WARP_IMAGE,
scale=(0.5, 0.75, 1.0, 1.25, 1.5),
flip=True,
)
# DataLoader
loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=CONFIG.BATCH_SIZE.TRAIN,
num_workers=CONFIG.NUM_WORKERS,
shuffle=True,
)
loader_iter = iter(loader)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
state_dict = torch.load(CONFIG.INIT_MODEL)
model.load_state_dict(state_dict, strict=False) # Skip "aspp" layer
model = nn.DataParallel(model)
model.to(device)
# Optimizer
optimizer = torch.optim.SGD(
# cf lr_mult and decay_mult in train.prototxt
params=[
{
"params": get_params(model.module, key="1x"),
"lr": CONFIG.LR,
"weight_decay": CONFIG.WEIGHT_DECAY,
},
{
"params": get_params(model.module, key="10x"),
"lr": 10 * CONFIG.LR,
"weight_decay": CONFIG.WEIGHT_DECAY,
},
{
"params": get_params(model.module, key="20x"),
"lr": 20 * CONFIG.LR,
"weight_decay": 0.0,
},
],
momentum=CONFIG.MOMENTUM,
)
# Loss definition
criterion = CrossEntropyLoss2d(ignore_index=CONFIG.IGNORE_LABEL)
criterion.to(device)
# TensorBoard Logger
writer = SummaryWriter(CONFIG.LOG_DIR)
loss_meter = MovingAverageValueMeter(20)
model.train()
model.module.scale.freeze_bn()
for iteration in tqdm(
range(1, CONFIG.ITER_MAX + 1),
total=CONFIG.ITER_MAX,
leave=False,
dynamic_ncols=True,
):
# Set a learning rate
poly_lr_scheduler(
optimizer=optimizer,
init_lr=CONFIG.LR,
iter=iteration - 1,
lr_decay_iter=CONFIG.LR_DECAY,
max_iter=CONFIG.ITER_MAX,
power=CONFIG.POLY_POWER,
)
# Clear gradients (ready to accumulate)
optimizer.zero_grad()
iter_loss = 0
for i in range(1, CONFIG.ITER_SIZE + 1):
try:
images, labels = next(loader_iter)
except:
loader_iter = iter(loader)
images, labels = next(loader_iter)
images = images.to(device)
labels = labels.to(device).unsqueeze(1).float()
# Propagate forward
logits = model(images)
# Loss
loss = 0
for logit in logits:
# Resize labels for {100%, 75%, 50%, Max} logits
labels_ = F.interpolate(labels,
logit.shape[2:],
mode="nearest")
labels_ = labels_.squeeze(1).long()
# Compute crossentropy loss
loss += criterion(logit, labels_)
# Backpropagate (just compute gradients wrt the loss)
loss /= float(CONFIG.ITER_SIZE)
loss.backward()
iter_loss += float(loss)
loss_meter.add(iter_loss)
# Update weights with accumulated gradients
optimizer.step()
# TensorBoard
if iteration % CONFIG.ITER_TB == 0:
writer.add_scalar("train_loss", loss_meter.value()[0], iteration)
for i, o in enumerate(optimizer.param_groups):
writer.add_scalar("train_lr_group{}".format(i), o["lr"],
iteration)
if False: # This produces a large log file
for name, param in model.named_parameters():
name = name.replace(".", "/")
writer.add_histogram(name, param, iteration, bins="auto")
if param.requires_grad:
writer.add_histogram(name + "/grad",
param.grad,
iteration,
bins="auto")
# Save a model
if iteration % CONFIG.ITER_SAVE == 0:
torch.save(
model.module.state_dict(),
osp.join(CONFIG.SAVE_DIR,
"checkpoint_{}.pth".format(iteration)),
)
# Save a model (short term)
if iteration % 100 == 0:
torch.save(
model.module.state_dict(),
osp.join(CONFIG.SAVE_DIR, "checkpoint_current.pth"),
)
torch.save(model.module.state_dict(),
osp.join(CONFIG.SAVE_DIR, "checkpoint_final.pth"))
def main(config, image_path, model_path, cuda, crf):
# Configuration
CONFIG = Dict(yaml.load(open(config)))
cuda = cuda and torch.cuda.is_available()
if cuda:
current_device = torch.cuda.current_device()
print('Running on', torch.cuda.get_device_name(current_device))
# Label list
with open(CONFIG.LABELS) as f:
classes = {}
for label in f:
label = label.rstrip().split('\t')
classes[int(label[0])] = label[1].split(',')[0]
# Load a model
state_dict = torch.load(model_path)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
model.load_state_dict(state_dict)
model.eval()
if cuda:
model.cuda()
image_size = (CONFIG.IMAGE.SIZE.TEST, ) * 2
# Image preprocessing
image = cv2.imread(image_path, cv2.IMREAD_COLOR).astype(float)
image = cv2.resize(image, image_size)
image_original = image.astype(np.uint8)
image -= np.array([
float(CONFIG.IMAGE.MEAN.B),
float(CONFIG.IMAGE.MEAN.G),
float(CONFIG.IMAGE.MEAN.R),
])
image = torch.from_numpy(image.transpose(2, 0, 1)).float().unsqueeze(0)
image = image.cuda() if cuda else image
# Inference
output = model(Variable(image, volatile=True))
output = F.upsample(output, size=image_size, mode='bilinear')
output = F.softmax(output, dim=1)
output = output.data.cpu().numpy()[0]
if crf:
output = dense_crf(image_original, output)
labelmap = np.argmax(output.transpose(1, 2, 0), axis=2)
labels = np.unique(labelmap)
# Show results
rows = np.floor(np.sqrt(len(labels) + 1))
cols = np.ceil((len(labels) + 1) / rows)
plt.figure(figsize=(10, 10))
ax = plt.subplot(rows, cols, 1)
ax.set_title('Input image')
ax.imshow(image_original[:, :, ::-1])
ax.set_xticks([])
ax.set_yticks([])
for i, label in enumerate(labels):
print('{0:3d}: {1}'.format(label, classes[label]))
mask = labelmap == label
ax = plt.subplot(rows, cols, i + 2)
ax.set_title(classes[label])
ax.imshow(image_original[:, :, ::-1])
ax.imshow(mask.astype(np.float32), alpha=0.5, cmap='viridis')
ax.set_xticks([])
ax.set_yticks([])
plt.tight_layout()
plt.show()
def main(config, model_path, cuda, crf):
cuda = cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda else "cpu")
if cuda:
current_device = torch.cuda.current_device()
print("Running on", torch.cuda.get_device_name(current_device))
else:
print("Running on CPU")
# Configuration
CONFIG = Dict(yaml.load(open(config)))
# Dataset 10k or 164k
dataset = get_dataset(CONFIG.DATASET)(
root=CONFIG.ROOT,
split=CONFIG.SPLIT.VAL,
base_size=CONFIG.IMAGE.SIZE.TEST,
mean=(CONFIG.IMAGE.MEAN.B, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.R),
warp=CONFIG.WARP_IMAGE,
scale=None,
flip=False,
)
# DataLoader
loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=CONFIG.BATCH_SIZE.TEST,
num_workers=CONFIG.NUM_WORKERS,
shuffle=False,
)
torch.set_grad_enabled(False)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)
model.load_state_dict(state_dict)
model = nn.DataParallel(model)
model.eval()
model.to(device)
preds, gts = [], []
for images, labels in tqdm(
loader, total=len(loader), leave=False, dynamic_ncols=True
):
# Image
images = images.to(device)
# Forward propagation
logits = model(images)
logits = F.interpolate(
logits, size=images.shape[2:], mode="bilinear", align_corners=True
)
probs = F.softmax(logits, dim=1)
probs = probs.data.cpu().numpy()
# Postprocessing
if crf:
pool = mp.Pool(mp.cpu_count())
images = images.data.cpu().numpy().astype(np.uint8).transpose(0, 2, 3, 1)
probs = pool.map(dense_crf_wrapper, zip(images, probs))
pool.close()
preds += list(np.argmax(probs, axis=1))
gts += list(labels.numpy())
score = scores(gts, preds, n_class=CONFIG.N_CLASSES)
with open(model_path.replace(".pth", ".json"), "w") as f:
json.dump(score, f, indent=4, sort_keys=True)
def train(config_path, cuda):
# Configuration
CONFIG = OmegaConf.load(config_path)
device = get_device(cuda)
torch.backends.cudnn.benchmark = True
# Dataset
dataset = get_dataset(CONFIG.DATASET.NAME)(
root=CONFIG.DATASET.ROOT,
split=CONFIG.DATASET.SPLIT.TRAIN,
ignore_label=CONFIG.DATASET.IGNORE_LABEL,
mean_bgr=(CONFIG.IMAGE.MEAN.B, CONFIG.IMAGE.MEAN.G,
CONFIG.IMAGE.MEAN.R),
augment=True,
base_size=CONFIG.IMAGE.SIZE.BASE,
crop_size=CONFIG.IMAGE.SIZE.TRAIN,
scales=CONFIG.DATASET.SCALES,
flip=True,
)
print(dataset)
# DataLoader
loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=CONFIG.SOLVER.BATCH_SIZE.TRAIN,
num_workers=CONFIG.DATALOADER.NUM_WORKERS,
shuffle=True,
)
loader_iter = iter(loader)
# Model check
print("Model:", CONFIG.MODEL.NAME)
assert (CONFIG.MODEL.NAME == "DeepLabV2_ResNet101_MSC"
), 'Currently support only "DeepLabV2_ResNet101_MSC"'
# Model setup
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.DATASET.N_CLASSES)
state_dict = torch.load(CONFIG.MODEL.INIT_MODEL)
print(" Init:", CONFIG.MODEL.INIT_MODEL)
for m in model.base.state_dict().keys():
if m not in state_dict.keys():
print(" Skip init:", m)
model.base.load_state_dict(state_dict, strict=False) # to skip ASPP
model = nn.DataParallel(model)
model.to(device)
# Loss definition
criterion = nn.CrossEntropyLoss(ignore_index=CONFIG.DATASET.IGNORE_LABEL)
criterion.to(device)
# Optimizer
optimizer = torch.optim.SGD(
# cf lr_mult and decay_mult in train.prototxt
params=[
{
"params": get_params(model.module, key="1x"),
"lr": CONFIG.SOLVER.LR,
"weight_decay": CONFIG.SOLVER.WEIGHT_DECAY,
},
{
"params": get_params(model.module, key="10x"),
"lr": 10 * CONFIG.SOLVER.LR,
"weight_decay": CONFIG.SOLVER.WEIGHT_DECAY,
},
{
"params": get_params(model.module, key="20x"),
"lr": 20 * CONFIG.SOLVER.LR,
"weight_decay": 0.0,
},
],
momentum=CONFIG.SOLVER.MOMENTUM,
)
# Learning rate scheduler
scheduler = PolynomialLR(
optimizer=optimizer,
step_size=CONFIG.SOLVER.LR_DECAY,
iter_max=CONFIG.SOLVER.ITER_MAX,
power=CONFIG.SOLVER.POLY_POWER,
)
# Setup loss logger
writer = SummaryWriter(
os.path.join(CONFIG.EXP.OUTPUT_DIR, "logs", CONFIG.EXP.ID))
average_loss = MovingAverageValueMeter(CONFIG.SOLVER.AVERAGE_LOSS)
# Path to save models
checkpoint_dir = os.path.join(
CONFIG.EXP.OUTPUT_DIR,
"models",
CONFIG.EXP.ID,
CONFIG.MODEL.NAME.lower(),
CONFIG.DATASET.SPLIT.TRAIN,
)
makedirs(checkpoint_dir)
print("Checkpoint dst:", checkpoint_dir)
model.train()
model.module.base.freeze_bn()
for iteration in tqdm(
range(1, CONFIG.SOLVER.ITER_MAX + 1),
total=CONFIG.SOLVER.ITER_MAX,
dynamic_ncols=True,
):
# Clear gradients (ready to accumulate)
optimizer.zero_grad()
loss = 0
for _ in range(CONFIG.SOLVER.ITER_SIZE):
try:
_, images, labels = next(loader_iter)
except:
loader_iter = iter(loader)
_, images, labels = next(loader_iter)
# Propagate forward
logits = model(images.to(device))
# Loss
iter_loss = 0
for logit in logits:
# Resize labels for {100%, 75%, 50%, Max} logits
_, _, H, W = logit.shape
labels_ = resize_labels(labels, size=(H, W))
iter_loss += criterion(logit, labels_.to(device))
# Propagate backward (just compute gradients)
iter_loss /= CONFIG.SOLVER.ITER_SIZE
iter_loss.backward()
loss += float(iter_loss)
average_loss.add(loss)
# Update weights with accumulated gradients
optimizer.step()
# Update learning rate
scheduler.step(epoch=iteration)
# TensorBoard
if iteration % CONFIG.SOLVER.ITER_TB == 0:
writer.add_scalar("loss/train", average_loss.value()[0], iteration)
for i, o in enumerate(optimizer.param_groups):
writer.add_scalar("lr/group_{}".format(i), o["lr"], iteration)
for i in range(torch.cuda.device_count()):
writer.add_scalar(
"gpu/device_{}/memory_cached".format(i),
torch.cuda.memory_cached(i) / 1024**3,
iteration,
)
if False:
for name, param in model.module.base.named_parameters():
name = name.replace(".", "/")
# Weight/gradient distribution
writer.add_histogram(name, param, iteration, bins="auto")
if param.requires_grad:
writer.add_histogram(name + "/grad",
param.grad,
iteration,
bins="auto")
# Save a model
if iteration % CONFIG.SOLVER.ITER_SAVE == 0:
torch.save(
model.module.state_dict(),
os.path.join(checkpoint_dir,
"checkpoint_{}.pth".format(iteration)),
)
torch.save(model.module.state_dict(),
os.path.join(checkpoint_dir, "checkpoint_final.pth"))
def main(config, model_path, cuda):
# Configuration
with open(config) as f:
CONFIG = yaml.load(f)
cuda = cuda and torch.cuda.is_available()
image_size = (CONFIG['IMAGE']['SIZE']['TEST'],
CONFIG['IMAGE']['SIZE']['TEST'])
n_classes = CONFIG['N_CLASSES']
# Dataset
dataset = get_dataset(CONFIG['DATASET'])(root=CONFIG['ROOT'],
split='test',
image_size=image_size,
scale=False,
flip=False,
preload=False)
# DataLoader
loader = torch.utils.data.DataLoader(dataset=dataset,
batch_size=CONFIG['BATCH_SIZE'],
num_workers=CONFIG['NUM_WORKERS'],
shuffle=False)
state_dict = torch.load(model_path,
map_location=lambda storage, loc: storage)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=n_classes)
model.load_state_dict(state_dict)
model.eval()
if cuda:
model.cuda()
targets, outputs = [], []
for data, target in tqdm(loader,
total=len(loader),
leave=False,
dynamic_ncols=True):
# Image
data = data.cuda() if cuda else data
data = Variable(data, volatile=True)
# Forward propagation
output = model(data)
output = F.upsample(output, size=image_size, mode='bilinear')
output = F.softmax(output, dim=1)
output = output.data.cpu().numpy()
crf_output = np.zeros(output.shape)
images = data.data.cpu().numpy().astype(np.uint8)
for i, (image, prob_map) in enumerate(zip(images, output)):
image = image.transpose(1, 2, 0)
crf_output[i] = dense_crf(image, prob_map)
output = crf_output
output = np.argmax(output, axis=1)
target = target.numpy()
for o, t in zip(output, target):
outputs.append(o)
targets.append(t)
score, class_iou = scores(targets, outputs, n_class=n_classes)
for k, v in score.items():
print k, v
score['Class IoU'] = {}
for i in range(n_classes):
score['Class IoU'][i] = class_iou[i]
with open('results.json', 'w') as f:
json.dump(score, f)
def main(config, model_path, cuda, crf):
# Configuration
CONFIG = Dict(yaml.load(open(config)))
cuda = cuda and torch.cuda.is_available()
if cuda:
current_device = torch.cuda.current_device()
print('Running on', torch.cuda.get_device_name(current_device))
image_size = (
CONFIG.IMAGE.SIZE.TEST,
CONFIG.IMAGE.SIZE.TEST,
)
# Dataset
dataset = CocoStuff10k(
root=CONFIG.ROOT,
split='test',
image_size=image_size,
scale=False,
flip=False,
preload=False,
)
# DataLoader
loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=CONFIG.BATCH_SIZE,
num_workers=CONFIG.NUM_WORKERS,
shuffle=False,
)
state_dict = torch.load(model_path,
map_location=lambda storage, loc: storage)
# Model
model = DeepLabV2_ResNet101_MSC(n_classes=CONFIG.N_CLASSES)
model.load_state_dict(state_dict)
model = nn.DataParallel(model)
model.eval()
if cuda:
model.cuda()
targets, outputs = [], []
for data, target in tqdm(
loader,
total=len(loader),
leave=False,
dynamic_ncols=True,
):
# Image
data = data.cuda() if cuda else data
data = Variable(data, volatile=True)
# Forward propagation
output = model(data)
output = F.upsample(output, size=image_size, mode='bilinear')
output = F.softmax(output, dim=1)
output = output.data.cpu().numpy()
# Postprocessing
if crf:
crf_output = np.zeros(output.shape)
images = data.data.cpu().numpy().astype(np.uint8)
for i, (image, prob_map) in enumerate(zip(images, output)):
image = image.transpose(1, 2, 0)
crf_output[i] = dense_crf(image, prob_map)
output = crf_output
output = np.argmax(output, axis=1)
target = target.numpy()
for o, t in zip(output, target):
outputs.append(o)
targets.append(t)
score, class_iou = scores(targets, outputs, n_class=CONFIG.N_CLASSES)
for k, v in score.items():
print(k, v)
score['Class IoU'] = {}
for i in range(CONFIG.N_CLASSES):
score['Class IoU'][i] = class_iou[i]
with open(model_path.replace('.pth', '.json'), 'w') as f:
json.dump(score, f, indent=4, sort_keys=True)
def main(config, embedding, model_path, run, imagedataset, local_rank, resnet,
bkg):
rank, world_size, device_id, device = setup(local_rank)
print("Local rank: {} Rank: {} World Size: {} Device_id: {} Device: {}".
format(local_rank, rank, world_size, device_id, device))
pth_extn = '.pth.tar'
# Configuration
CONFIG = Dict(yaml.load(open(config)))
datadir = os.path.join('data/datasets', imagedataset)
print("Split dir: ", datadir)
savedir = osp.dirname(model_path)
epoch = re.findall("checkpoint_(.*)\." + pth_extn[1:],
osp.basename(model_path))[-1]
if run == 'zlss' or run == 'flss':
val = np.load(datadir + '/split/test_list.npy')
visible_classes = np.load(datadir + '/split/novel_cls.npy')
if bkg:
visible_classes = np.asarray(np.concatenate(
[np.array([0]), visible_classes]),
dtype=int)
elif run == 'gzlss' or run == 'gflss':
val = np.load(datadir + '/split/test_list.npy')
vals_cls = np.asarray(np.concatenate([
np.load(datadir + '/split/seen_cls.npy'),
np.load(datadir + '/split/val_cls.npy')
]),
dtype=int)
if bkg:
vals_cls = np.asarray(np.concatenate([np.array([0]), vals_cls]),
dtype=int)
valu_cls = np.load(datadir + '/split/novel_cls.npy')
visible_classes = np.concatenate([vals_cls, valu_cls])
else:
print("invalid run ", run)
sys.exit()
cls_map = np.array([255] * 256)
for i, n in enumerate(visible_classes):
cls_map[n] = i
if run == 'gzlss' or run == 'gflss':
novel_cls_map = np.array([255] * 256)
for i, n in enumerate(list(valu_cls)):
novel_cls_map[cls_map[n]] = i
seen_cls_map = np.array([255] * 256)
for i, n in enumerate(list(vals_cls)):
seen_cls_map[cls_map[n]] = i
all_labels = np.genfromtxt(datadir + '/labels_2.txt',
delimiter='\t',
usecols=1,
dtype='str')
print("Visible Classes: ", visible_classes)
# Dataset
dataset = get_dataset(CONFIG.DATASET)(
train=None,
test=val,
root=CONFIG.ROOT,
split=CONFIG.SPLIT.TEST,
base_size=CONFIG.IMAGE.SIZE.TEST,
mean=(CONFIG.IMAGE.MEAN.B, CONFIG.IMAGE.MEAN.G, CONFIG.IMAGE.MEAN.R),
warp=CONFIG.WARP_IMAGE,
scale=None,
flip=False,
)
random.seed(42)
if embedding == 'word2vec':
class_emb = pickle.load(
open(datadir + '/word_vectors/word2vec.pkl', "rb"))
elif embedding == 'fasttext':
class_emb = pickle.load(
open(datadir + '/word_vectors/fasttext.pkl', "rb"))
elif embedding == 'fastnvec':
class_emb = np.concatenate([
pickle.load(open(datadir + '/word_vectors/fasttext.pkl', "rb")),
pickle.load(open(datadir + '/word_vectors/word2vec.pkl', "rb"))
],
axis=1)
else:
print("invalid emb ", embedding)
sys.exit()
class_emb = class_emb[visible_classes]
class_emb = F.normalize(torch.tensor(class_emb), p=2, dim=1).cuda()
print("Embedding dim: ", class_emb.shape[1])
print("# Visible Classes: ", class_emb.shape[0])
# DataLoader
loader = torch.utils.data.DataLoader(dataset=dataset,
batch_size=CONFIG.BATCH_SIZE.TEST,
num_workers=CONFIG.NUM_WORKERS,
shuffle=False,
sampler=DistributedSampler(
dataset,
num_replicas=world_size,
rank=rank,
shuffle=False),
pin_memory=True,
drop_last=True)
torch.set_grad_enabled(False)
# Model
model = DeepLabV2_ResNet101_MSC(class_emb.shape[1],
class_emb,
resnet=resnet)
state_dict = torch.load(model_path, map_location='cpu')
model = DistributedDataParallel(model.to(device), device_ids=[rank])
new_state_dict = OrderedDict()
if resnet == 'spnet':
for k, v in state_dict['state_dict'].items():
name = k.replace("scale", "base") # 'scale'->base
name = name.replace("stages.", "")
new_state_dict[name] = v
else:
new_state_dict = state_dict['state_dict']
model.load_state_dict(new_state_dict)
del state_dict
model.eval()
targets, outputs = [], []
loader_iter = iter(loader)
iterations = len(loader_iter)
print("Iterations: {}".format(iterations))
pbar = tqdm(loader,
total=iterations,
leave=False,
dynamic_ncols=True,
position=rank)
for iteration in pbar:
data, target, img_id = next(loader_iter)
# Image
data = data.to(device)
# Forward propagation
output = model(data)
output = F.interpolate(output,
size=data.shape[2:],
mode="bilinear",
align_corners=False)
output = F.softmax(output, dim=1)
target = cls_map[target.numpy()]
remote_target = torch.tensor(target).to(device)
if rank == 0:
remote_target = torch.zeros_like(remote_target).to(device)
output = torch.argmax(output, dim=1).cpu().numpy()
remote_output = torch.tensor(output).to(device)
if rank == 0:
remote_output = torch.zeros_like(remote_output).to(device)
for o, t in zip(output, target):
outputs.append(o)
targets.append(t)
torch.distributed.reduce(remote_output, dst=0)
torch.distributed.reduce(remote_target, dst=0)
torch.distributed.barrier()
if rank == 0:
remote_output = remote_output.cpu().numpy()
remote_target = remote_target.cpu().numpy()
for o, t in zip(remote_output, remote_target):
outputs.append(o)
targets.append(t)
if rank == 0:
if run == 'gzlss' or run == 'gflss':
score, class_iou = scores_gzsl(targets,
outputs,
n_class=len(visible_classes),
seen_cls=cls_map[vals_cls],
unseen_cls=cls_map[valu_cls])
else:
score, class_iou = scores(targets,
outputs,
n_class=len(visible_classes))
for k, v in score.items():
print(k, v)
score["Class IoU"] = {}
for i in range(len(visible_classes)):
score["Class IoU"][all_labels[visible_classes[i]]] = class_iou[i]
name = ""
name = model_path.replace(pth_extn, "_" + run + ".json")
if bkg == True:
with open(name.replace('.json', '_bkg.json'), "w") as f:
json.dump(score, f, indent=4, sort_keys=True)
else:
with open(name, "w") as f:
json.dump(score, f, indent=4, sort_keys=True)
print(score["Class IoU"])
return
