def test(config):
fields = rstdt_fields()
dataset = Dataset(config, fields)
test_iter = dataset.load_test()
# model = SpanBasedParser.load_model(config.model_path[0], config, fields)
# model = EnsembleParser.load_model(config.model_path, config, fields)
model = HierarchicalParser.load_model(config.model_path, config, fields)
# scores = Trainer.valid(model, test_iter)
# print(scores)
doc_ids = []
pred_trees = []
for batch in test_iter:
batch.tree = None
with torch.no_grad():
output_dict = model(batch)
doc_ids.extend(batch.doc_id)
pred_trees.extend(output_dict['tree'])
config.output_dir.mkdir(parents=True, exist_ok=True)
pred_trees = [Tree.fromstring(tree) for tree in pred_trees]
for doc_id, tree in zip(doc_ids, pred_trees):
tree_path = config.output_dir / '{}.tree'.format(doc_id)
with open(tree_path, 'w') as f:
print(tree, file=f)
return
def main():
params = parse_args()
tf.random.set_seed(params.seed)
tf.get_logger().setLevel(logging.ERROR)
params = setup_horovod(params)
set_flags(params)
model_dir = prepare_model_dir(params)
logger = get_logger(params)
dataset = Dataset(data_dir=params.data_dir,
batch_size=params.batch_size,
fold_idx=params.fold,
n_folds=params.num_folds,
params=params,
seed=params.seed)
estimator = build_estimator(params, model_dir)
if params.tensorboard_logging and (params.worker_id == 0
or params.log_all_workers):
from TensorFlow.common.tb_utils import write_hparams_v1
write_hparams_v1(params.log_dir, vars(params))
if not params.benchmark:
params.max_steps = params.max_steps // params.num_workers
if 'train' in params.exec_mode:
with dump_callback(params.dump_config):
training_hooks = get_hooks(params, logger)
dataset_fn = dataset.synth_train_fn if params.synth_data else dataset.train_fn
estimator.train(input_fn=dataset_fn,
steps=params.max_steps,
hooks=training_hooks)
if 'evaluate' in params.exec_mode:
result = estimator.evaluate(input_fn=dataset.eval_fn,
steps=dataset.eval_size)
data = parse_evaluation_results(result)
if params.worker_id == 0:
logger.log(step=(), data=data)
if 'predict' == params.exec_mode:
inference_hooks = get_hooks(params, logger)
if params.worker_id == 0:
count = 1 if not params.benchmark else 2 * params.warmup_steps * params.batch_size // dataset.test_size
predictions = estimator.predict(input_fn=lambda: dataset.test_fn(
count=count, drop_remainder=params.benchmark),
hooks=inference_hooks)
for idx, p in enumerate(predictions):
volume = p['predictions']
if not params.benchmark:
np.save(
os.path.join(params.model_dir,
"vol_{}.npy".format(idx)), volume)
def test(args, cfg):
cudnn.benchmark = True
device = torch.device(
'cuda' if args.gpu=="0" else 'cpu')
if args.weights:
if os.path.isdir(args.weights):
weights = find_best_weights(args.weights, cfg)
weights = [weights[0], weights[1]]
print(weights)
models = [restore_model(get_model(), weight) for weight in weights]
else:
models = [restore_model(get_model(), args.weights)]
else:
models = [get_model(freeze=True) for _ in range(cfg.kfolds)]
for model in models:
model.eval()
with open('csv/res_{}.csv'.format(args.weights.split('/')[1]), 'w') as csvfile:
writer = csv.writer(csvfile)
writer.writerow(['id', 'rle_mask'])
test_set = Dataset(args.datadir, config, phase='test', fold_idx=None, mean=model.mean, std=model.std)
test_loader = DataLoader(test_set, batch_size=1, shuffle=False)
tta = pad_and_crop(224)
for image, depth, image_id, zero_flag in tqdm.tqdm(test_loader):
if zero_flag.numpy()[0]:
mask = np.zeros([cfg.h, cfg.w, 1], dtype=np.uint8)
else:
image = image.to(device)
depth = depth.to(device)
image = torch.cat([image, image.flip(3)], dim=0)
depth = torch.cat([depth, depth], dim=0)
logits = torch.zeros(128, 128, dtype=torch.float32).to(device)
for model in models:
model.to(device)
y_pred = model(image, depth)
logits += torch.cat([y_pred[0].unsqueeze(0), y_pred[1].unsqueeze(0).flip(3)], dim=0).mean(dim=0).squeeze(0)
mask = ((logits / len(models))>0.5).byte().cpu().numpy()
if cfg.use_tta:
mask = pad_and_crop(224, 202)(image=mask)['image']
mask = cv2.resize(mask, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_AREA)
else:
mask = crop(mask, cfg.h, cfg.w)[:, :, np.newaxis]
mask = remove_small_objects(label(mask), 10).astype(np.uint8)
mask[mask > 1] = 1
rles = list(mask_to_rles(mask))
for rle in rles:
writer.writerow(
[image_id[0], ' '.join(str(y) for y in rle)])
if len(rles) == 0:
writer.writerow([image_id[0], ''])
def train(config):
fields = rstdt_fields()
dataset = Dataset(config, fields)
train_iter, valid_iter = dataset.load_train_valid()
model = SpanBasedParser.build_model(config, fields)
print(model)
optimizer = {
'adam': optim.Adam,
'sgd': optim.SGD
}[config.optimizer](model.parameters(),
lr=config.lr,
weight_decay=config.weight_decay)
scheduler = optim.lr_scheduler.ExponentialLR(optimizer, config.lr_decay)
trainer = Trainer(config, model, optimizer, scheduler, train_iter,
valid_iter, fields)
trainer.run()
return
def main():
""" Starting point of the application """
hvd.init()
set_flags()
params = PARSER.parse_args()
logger = get_logger(params)
dataset = Dataset(data_dir=params.data_dir,
batch_size=params.batch_size,
fold_idx=params.fold,
n_folds=params.num_folds,
input_shape=params.input_shape,
params=params)
estimator = build_estimator(params=params, model_fn=unet_3d)
hooks = get_hooks(params, logger)
if 'train' in params.exec_mode:
max_steps = params.max_steps // (1 if params.benchmark else hvd.size())
estimator.train(input_fn=dataset.train_fn,
steps=max_steps,
hooks=hooks)
if 'evaluate' in params.exec_mode:
result = estimator.evaluate(input_fn=dataset.eval_fn,
steps=dataset.eval_size)
_ = parse_evaluation_results(result, logger)
if params.exec_mode == 'predict':
if hvd.rank() == 0:
predictions = estimator.predict(input_fn=dataset.test_fn,
hooks=hooks)
for idx, pred in enumerate(predictions):
volume = pred['predictions']
if not params.benchmark:
np.save(
os.path.join(params.model_dir,
"vol_{}.npy".format(idx)), volume)
def train(args, cfg):
#cudnn.benchmark = True
device = torch.device('cuda' if args.gpu == "0" else 'cpu')
if args.weights:
#import ipdb; ipdb.set_trace()
weights = find_best_weights(args.weights, cfg)
#end = cfg.start_fold if cfg.start_cycle == 0 else cfg.start_fold + 1
models = [
restore_model(get_model(), weight) if weight else get_model(
freeze=True) for weight in weights
]
else:
models = [get_model(freeze=True) for _ in range(cfg.kfolds)]
lr = cfg.lr
for fold in range(cfg.start_fold, cfg.kfolds):
#if fold in [0, 1, 2]:
# continue
best_metric = float(weights[fold].split('-')
[3]) if args.weights and weights[fold] else -np.inf
best_tta_metric = -np.inf
start_cycle = cfg.start_cycle if fold == cfg.start_fold else 0
for cycle in range(start_cycle, cfg.num_cycles):
#import ipdb; ipdb.set_trace()
lr = cfg.start_lr if fold == cfg.start_fold and cycle == cfg.start_cycle else lr
alpha_zero = lr
if cycle != 0 and best_metric < 0.83:
break
model = models[fold]
print(sum(p.numel() for p in model.parameters()))
print("Best metric is {}".format(best_metric))
optimizer = torch.optim.Adam(
model.parameters(),
lr=lr) if cycle == 0 else torch.optim.RMSprop(
model.parameters(), lr=lr)
lr_scheduler = ReduceLROnPlateau(optimizer,
'max',
factor=cfg.lr_mult,
patience=16,
verbose=True)
metric_fn = get_metric()
model.to(device)
train_set = Dataset(args.datadir,
config,
phase='train',
fold_idx=fold,
mean=model.mean,
std=model.std)
val_set = Dataset(args.datadir,
config,
phase='val',
fold_idx=fold,
mean=model.mean,
std=model.std)
val_set_tta = Dataset(args.datadir,
config,
phase='val',
fold_idx=fold,
mean=model.mean,
std=model.std,
use_tta=True)
train_sampler = BatchSampler(Sampler(train_set), cfg.batch_size)
train_loader = DataLoader(train_set,
batch_sampler=train_sampler,
num_workers=8)
val_loader = DataLoader(val_set,
batch_size=cfg.batch_size,
shuffle=True)
val_tta_loader = DataLoader(val_set_tta,
batch_size=cfg.batch_size,
shuffle=True)
def train_step(loader):
train_loss = []
train_metric = []
for image, depth, mask in tqdm.tqdm(loader):
image = image.type(torch.float).to(device)
mask = mask.to(device)
depth = depth.to(device)
y_pred = model(image, depth)
loss = loss_fn(y_pred, mask)
metric = metric_fn(y_pred, mask)
train_metric.append(metric.item())
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_loss.append(loss.item())
return train_loss, train_metric
def val_step(loader):
val_loss = []
val_metric = []
with torch.no_grad():
for image, depth, mask in tqdm.tqdm(loader):
image = image.to(device)
mask = mask.to(device)
depth = depth.to(device)
y_pred = model(image, depth)
loss = loss_fn(y_pred, mask)
metric = metric_fn((y_pred > 0.5).byte(), mask.byte())
val_loss.append(loss.item())
val_metric.append(metric.item())
return val_loss, val_metric
if cycle == 0 and not args.weights:
loss_fn = get_loss(0, freeze_epochs=True)
for epoch in range(cfg.frozen_epochs):
print(train_set.pad_size, " ", train_set.crop_size)
train_loss, train_metric = train_step(train_loader)
print("Epoch: %d, Train Loss: %.4f, Train Metric: %.4f," %
(epoch, np.mean(train_loss), np.mean(train_metric)))
for param in model.parameters():
param.requires_grad = True
for param_group in optimizer.param_groups:
param_group['lr'] = lr
num_epochs = cfg.num_epochs if cycle > 0 else cfg.start_epochs
start_epoch = cfg.start_epoch if cycle == start_cycle and fold == cfg.start_fold else 0
for epoch in range(start_epoch, num_epochs):
loss_fn = get_loss(cycle)
print(loss_fn)
print(train_set.pad_size, " ", train_set.crop_size)
train_loss, train_metric = train_step(train_loader)
val_loss, val_metric = val_step(val_loader)
val_tta_loss, val_tta_metric = val_step(val_tta_loader)
print(
"Epoch: %d, Train Loss: %.4f, Train Metric: %.4f, Val Loss: %.4f, Val Metric: %.4f, Val TTA Loss: %.4f, Val TTA Metric: %.4f"
% (
epoch,
np.mean(train_loss),
np.mean(train_metric),
np.mean(val_loss),
np.mean(val_metric),
np.mean(val_tta_loss),
np.mean(val_tta_metric),
))
lr = update_lr(optimizer,
cfg,
lr,
epoch,
alpha_zero=alpha_zero,
verbose=1)
lr_scheduler.step(np.mean(val_metric))
new_lr = optimizer.param_groups[0]['lr']
if abs(new_lr - lr) > 1e-7:
alpha_zero = max(lr * cfg.lr_mult, 1.6e-5)
best_metric, best_tta_metric = save_checkpoint(
model,
np.mean(val_metric),
best_metric,
np.mean(val_tta_metric),
best_tta_metric,
fold,
epoch,
verbose=1)
def main():
tf.get_logger().setLevel(logging.ERROR)
hvd.init()
params = PARSER.parse_args()
model_dir = prepare_model_dir(params)
logger = get_logger(params)
dataset = Dataset(data_dir=params.data_dir,
batch_size=params.batch_size,
fold_idx=params.fold,
n_folds=params.num_folds,
params=params)
estimator = build_estimator(params=params, model_dir=model_dir)
max_steps = params.max_steps // (1 if params.benchmark else hvd.size())
if 'train' in params.exec_mode:
training_hooks = get_hooks(params, logger)
estimator.train(input_fn=dataset.train_fn,
steps=max_steps,
hooks=training_hooks)
if 'evaluate' in params.exec_mode:
result = estimator.evaluate(input_fn=dataset.eval_fn,
steps=dataset.eval_size)
data = parse_evaluation_results(result)
if hvd.rank() == 0:
logger.log(step=(), data=data)
if 'predict' == params.exec_mode:
inference_hooks = get_hooks(params, logger)
if hvd.rank() == 0:
count = 1 if not params.benchmark else 2 * params.warmup_steps * params.batch_size // dataset.test_size
predictions = estimator.predict(input_fn=lambda: dataset.test_fn(
count=count, drop_remainder=params.benchmark),
hooks=inference_hooks)
for idx, p in enumerate(predictions):
volume = p['predictions']
if not params.benchmark:
np.save(
os.path.join(params.model_dir,
"vol_{}.npy".format(idx)), volume)
if 'debug_train' == params.exec_mode:
hooks = [hvd.BroadcastGlobalVariablesHook(0)]
if hvd.rank() == 0:
hooks += [
TrainingHook(log_every=params.log_every,
logger=logger,
tensor_names=['total_loss_ref:0']),
ProfilingHook(warmup_steps=params.warmup_steps,
global_batch_size=hvd.size() * params.batch_size,
logger=logger,
mode='train')
]
estimator.train(input_fn=dataset.synth_train_fn,
steps=max_steps,
hooks=hooks)
if 'debug_predict' == params.exec_mode:
if hvd.rank() == 0:
hooks = [
ProfilingHook(warmup_steps=params.warmup_steps,
global_batch_size=params.batch_size,
logger=logger,
mode='inference')
]
count = 2 * params.warmup_steps
predictions = estimator.predict(
input_fn=lambda: dataset.synth_predict_fn(count=count),
hooks=hooks)
for p in predictions:
_ = p['predictions']