def main(cfg, gpu):
torch.cuda.set_device(gpu)
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = ModelBuilder.build_decoder(
arch=cfg.MODEL.arch_decoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
num_class=cfg.DATASET.num_class,
weights=cfg.MODEL.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
# Dataset and Loader
dataset_val = ValDataset(cfg.DATASET.root_dataset, cfg.DATASET.list_val,
cfg.DATASET)
loader_val = torch.utils.data.DataLoader(dataset_val,
batch_size=cfg.VAL.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
data_set_rec = ValDataset(cfg.DATASET.root_dataset,
cfg.DATASET.list_val,
cfg.DATASET,
rec_dataset=cfg.DATASET.rec_dataset)
loader_rec = torch.utils.data.DataLoader(
data_set_rec,
batch_size=cfg.VAL.batch_size, # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
# Main loop
evaluate(segmentation_module, loader_val, loader_rec, cfg, gpu)
print('Evaluation Done!')
def worker(args, dev_id, start_idx, end_idx, result_queue):
torch.cuda.set_device(dev_id)
# Dataset and Loader
dataset_val = ValDataset(
broden_dataset.record_list['validation'], args,
max_sample=args.num_val, start_idx=start_idx,
end_idx=end_idx)
loader_val = torchdata.DataLoader(
dataset_val,
batch_size=args.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=2)
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(
arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(
arch=args.arch_decoder,
fc_dim=args.fc_dim,
nr_classes=args.nr_classes,
weights=args.weights_decoder,
use_softmax=True)
segmentation_module = SegmentationModule(net_encoder, net_decoder)
segmentation_module.cuda()
# Main loop
evaluate(segmentation_module, loader_val, args, dev_id, result_queue)
def main(args):
# import network architecture
builder = ModelBuilder()
model = builder.build_net(arch=args.id,
num_input=args.num_input,
num_classes=args.num_classes,
num_branches=args.num_branches,
padding_list=args.padding_list,
dilation_list=args.dilation_list)
model = torch.nn.DataParallel(model, device_ids=list(range(
args.num_gpus))).cuda()
cudnn.benchmark = True
if args.resume:
if os.path.isfile(args.resume):
print("=> Loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
state_dict = checkpoint['state_dict']
model.load_state_dict(state_dict)
print("=> Loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
else:
raise Exception("=> No checkpoint found at '{}'".format(
args.resume))
tf = ValDataset(test_dir, args)
test_loader = DataLoader(tf,
batch_size=args.batch_size,
shuffle=args.shuffle,
num_workers=args.num_workers,
pin_memory=False)
test(test_loader, model, args)
def main(args):
if args.random:
print("Randomly select image from ValDataset")
val_dataset = ValDataset()
rand_idx = random.randint(0, len(val_dataset) - 1)
img, label = val_dataset[rand_idx]
else:
print("Loading image from {}".format(args.image))
img = Image.open(args.image)
img = transforms.functional.resize(img, (HEIGHT, WIDTH))
img = transforms.functional.to_tensor(img)
img = img.unsqueeze(0).to(device)
ckpt = torch.load(args.checkpoint,
map_location=lambda storage, loc: storage)
model = LPRModel().to(device)
model.load_state_dict(ckpt)
model.eval()
print('weight has been loaded')
predict = model(img)
predict = torch.argmax(predict, axis=-1).cpu()
if args.random:
print('label:', decode(label.numpy()))
print('predictions:', decode(predict[0].numpy()))
def worker(args, gpu_id, start_idx, end_idx, result_queue):
torch.cuda.set_device(gpu_id)
# Dataset and Loader
dataset_val = ValDataset(args.list_val,
args,
max_sample=args.num_val,
start_idx=start_idx,
end_idx=end_idx)
loader_val = torchdata.DataLoader(dataset_val,
batch_size=args.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=2)
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
num_class=args.num_class,
weights=args.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
segmentation_module.cuda()
# Main loop
evaluate(segmentation_module, loader_val, args, gpu_id, result_queue)
def main(args):
torch.cuda.set_device(args.gpu_id)
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
weights=args.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
# Dataset and Loader
dataset_val = ValDataset(args.list_val, args, max_sample=args.num_val)
loader_val = torchdata.DataLoader(dataset_val,
batch_size=args.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
# Main loop
evaluate(segmentation_module, loader_val, args)
print('Evaluation Done!')
def worker(cfg, gpu_id, start_idx, end_idx, result_queue):
torch.cuda.set_device(gpu_id)
# Dataset and Loader
dataset_val = ValDataset(cfg.DATASET.root_dataset,
cfg.DATASET.list_val,
cfg.DATASET,
start_idx=start_idx,
end_idx=end_idx)
loader_val = torch.utils.data.DataLoader(dataset_val,
batch_size=cfg.VAL.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=2)
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = ModelBuilder.build_decoder(
arch=cfg.MODEL.arch_decoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
num_class=cfg.DATASET.num_class,
weights=cfg.MODEL.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
segmentation_module.cuda()
# Main loop
evaluate(segmentation_module, loader_val, cfg, gpu_id, result_queue)
def __init__(self, split='trainaug', dtype=None, val=True):
torch.manual_seed(66)
torch.cuda.manual_seed_all(66)
self.log_dir = osp.join(settings.LOG_DIR, settings.EXP_NAME)
self.model_dir = osp.join(settings.MODEL_DIR, settings.EXP_NAME)
ensure_dir(self.log_dir)
ensure_dir(self.model_dir)
logger.info('set log dir as %s' % self.log_dir)
logger.info('set model dir as %s' % self.model_dir)
self.step = 1
self.best_mIoU = 0
self.writer = SummaryWriter(self.log_dir)
self.split = split
train_set = TrainDataset(split=split)
self.train_loader = DataLoader(train_set,
batch_size=settings.TRAIN_BATCH_SIZE,
pin_memory=True,
num_workers=settings.NUM_WORKERS,
shuffle=True,
drop_last=True)
val_set = ValDataset(split='val')
self.val_loader = DataLoader(val_set,
batch_size=1,
shuffle=False,
num_workers=settings.NUM_WORKERS,
drop_last=False)
self.net = HamNet(settings.N_CLASSES, settings.N_LAYERS).cuda()
params_count(self.net)
self.opt = SGD(params=[{
'params': get_params(self.net, key='1x'),
'lr': 1 * settings.LR,
'weight_decay': settings.WEIGHT_DECAY,
}, {
'params': get_params(self.net, key='1y'),
'lr': 1 * settings.LR,
'weight_decay': 0,
}, {
'params': get_params(self.net, key='2x'),
'lr': 2 * settings.LR,
'weight_decay': 0.0,
}],
momentum=settings.LR_MOM)
self.net = DataParallel(self.net)
patch_replication_callback(self.net)
def __init__(self, dt_split):
torch.cuda.set_device(settings.DEVICE)
self.log_dir = settings.LOG_DIR
self.model_dir = settings.MODEL_DIR
self.net = EANet(settings.N_CLASSES, settings.N_LAYERS).cuda()
self.net = DataParallel(self.net, device_ids=[settings.DEVICE])
dataset = ValDataset(split=dt_split)
self.dataloader = DataLoader(dataset, batch_size=1, shuffle=False,
num_workers=2, drop_last=False)
self.hist = 0
def __init__(self, dt_split):
self.model_dir = osp.join(settings.MODEL_DIR, settings.EXP_NAME)
self.net = HamNet(settings.N_CLASSES, settings.N_LAYERS).cuda()
self.net = DataParallel(self.net)
dataset = ValDataset(split=dt_split)
self.dataloader = DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=2,
drop_last=False)
self.hist = 0
def prepare_data(self):
"""prepare_data
load dataset
"""
d = self.hparams['data']
self.train_ds = Dataset(
'./data/articles.json', './data/users_list.json', self.w2v, maxlen=self.hparams['data']['maxlen'], pos_num=d['pos_k'], neg_k=d['neg_k'])
self.val_ds = ValDataset(
50, './data/articles.json', './data/users_list.json', self.w2v)
tmp = [t.unsqueeze(0) for t in self.train_ds[0]]
self.logger.experiment.add_graph(self.model, tmp)
def main(cfg, gpu):
torch.cuda.set_device(gpu)
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = ModelBuilder.build_decoder(
arch=cfg.MODEL.arch_decoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
num_class=cfg.DATASET.num_class,
weights=cfg.MODEL.weights_decoder,
use_softmax=True)
# use_softmax = True
crit = nn.NLLLoss(ignore_index=-1)
if cfg.MODEL.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit,
cfg.TRAIN.deep_sup_scale)
else:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit)
# Dataset and Loader
dataset_val = ValDataset(cfg.DATASET.root_dataset, cfg.DATASET.list_val,
cfg.DATASET)
loader_val = torch.utils.data.DataLoader(dataset_val,
batch_size=cfg.VAL.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
# Main loop
evaluate(segmentation_module, loader_val, cfg, gpu)
print('Evaluation Done!')
def main(args):
torch.cuda.set_device(args.gpu)
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
num_class=args.num_class,
weights=args.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
# Dataset and Loader
#if len(args.test_imgs) == 1 and os.path.isdir(args.test_imgs[0]):
# test_imgs = find_recursive(args.test_imgs[0])
#else:
# test_imgs = args.test_imgs
#list_test = [{'fpath_img': x} for x in test_imgs]
dataset_test = ValDataset(args.test_imgs, args, max_sample=args.num_val)
# test_set = dl.loadTest()
loader_test = torchdata.DataLoader(dataset_test,
batch_size=args.batch_size,
shuffle=False,
num_workers=1,
drop_last=True)
segmentation_module.cuda()
# Main loop
test(segmentation_module, loader_test, args)
print('Inference done!')
def prepare_data(self):
"""prepare_data
load dataset
"""
print('Prepating train data...', end=' ')
self.train_ds = Dataset(self.hparams['all_items'],
self.hparams['train_data'],
self.w2v,
maxlen=self.hparams['title_len'],
npratio=self.hparams['train']['npratio'],
his_len=self.hparams['his_len'])
self.val_ds = ValDataset(self.hparams['all_items'],
self.hparams['val_data'],
self.hparams['val_items'],
self.w2v,
maxlen=self.hparams['title_len'],
his_len=self.hparams['his_len'],
pos=self.hparams['val']['pos'],
neg=self.hparams['val']['neg'],
total=self.hparams['val']['total'])
print('Done')
def main(**args):
r"""Performs the main training loop
"""
gray_mode = args['gray'] # gray mode indicator
C = 1 if gray_mode else 3 # number of color channels
# Load dataset
print('> Loading datasets ...')
dataset_val = ValDataset(valsetdir=args['valset_dir'],
gray_mode=gray_mode) # for grayscale/color video
# dataset_val = ValDataset(valsetdir=args['valset_dir'], gray_mode=False) # for color videos only
loader_train = train_dali_loader(batch_size=args['batch_size'],\
file_root=args['trainset_dir'],\
sequence_length=args['temp_patch_size'],\
crop_size=args['patch_size'],\
epoch_size=args['max_number_patches'],\
random_shuffle=True,\
temp_stride=3,\
gray_mode=gray_mode)
num_minibatches = int(args['max_number_patches'] // args['batch_size'])
ctrl_fr_idx = (args['temp_patch_size'] - 1) // 2
print("\t# of training samples: %d\n" % int(args['max_number_patches']))
# Init loggers
writer, logger = init_logging(args)
# Define GPU devices
device_ids = [0]
torch.backends.cudnn.benchmark = True # CUDNN optimization
# Create model
model = FastDVDnet(num_color_channels=C)
model = model.cuda()
# Define loss
criterion = nn.MSELoss(reduction='sum')
criterion.cuda()
# Optimizer
optimizer = optim.Adam(model.parameters(), lr=args['lr'])
# [AMP initialization] automated half-precision training
if args['fp16']:
model, optimizer = amp.initialize(model,
optimizer,
opt_level=args['amp_opt_level'])
# model = nn.DataParallel(model, device_ids=device_ids).cuda()
model = nn.DataParallel(model)
# Resume training or start anew
start_epoch, training_params = resume_training(args, model, optimizer)
# Training
start_time = time.time()
for epoch in range(start_epoch, args['epochs']):
# Set learning rate
current_lr, reset_orthog = lr_scheduler(epoch, args)
if reset_orthog:
training_params['no_orthog'] = True
# set learning rate in optimizer
for param_group in optimizer.param_groups:
param_group["lr"] = current_lr
print('\nlearning rate %f' % current_lr)
# train
for i, data in enumerate(loader_train, 0):
# Pre-training step
model.train()
# When optimizer = optim.Optimizer(net.parameters()) we only zero the optim's grads
optimizer.zero_grad()
# convert inp to [N, num_frames*C. H, W] in [0., 1.] from [N, num_frames, C. H, W] in [0., 255.]
# extract ground truth (central frame)
img_train, gt_train = normalize_augment(data[0]['data'],
ctrl_fr_idx, gray_mode)
N, _, H, W = img_train.size()
# std dev of each sequence
stdn = torch.empty(
(N, 1, 1, 1)).cuda().uniform_(args['noise_ival'][0],
to=args['noise_ival'][1])
# draw noise samples from std dev tensor
noise = torch.zeros_like(img_train)
noise = torch.normal(mean=noise, std=stdn.expand_as(noise))
#define noisy input
imgn_train = img_train + noise
# Send tensors to GPU
gt_train = gt_train.cuda(non_blocking=True)
imgn_train = imgn_train.cuda(non_blocking=True)
noise = noise.cuda(non_blocking=True)
noise_map = stdn.expand(
(N, 1, H, W)).cuda(non_blocking=True) # one channel per image
# Evaluate model and optimize it
out_train = model(imgn_train, noise_map)
# Compute loss
loss = criterion(gt_train, out_train) / (N * 2)
# [AMP scale loss to avoid overflow of float16] automated mixed precision training
if args['fp16']:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
optimizer.step()
# Results
if training_params['step'] % args['save_every'] == 0:
# Apply regularization by orthogonalizing filters
if not training_params['no_orthog']:
model.apply(svd_orthogonalization)
# Compute training PSNR
log_train_psnr(out_train, \
gt_train, \
loss, \
writer, \
epoch, \
i, \
num_minibatches, \
training_params)
# update step counter
training_params['step'] += 1
# save model and checkpoint
training_params['start_epoch'] = epoch + 1
save_model_checkpoint(model, args, optimizer, training_params, epoch)
# Call to model.eval() to correctly set the BN layers before inference
model.eval()
# Validation and log images
validate_and_log(
model_temp=model, \
dataset_val=dataset_val, \
valnoisestd=args['val_noiseL'], \
temp_psz=args['temp_patch_size'], \
writer=writer, \
epoch=epoch, \
lr=current_lr, \
logger=logger, \
trainimg=img_train, \
gray_mode=gray_mode
)
# Print elapsed time
elapsed_time = time.time() - start_time
print('Elapsed time {}'.format(
time.strftime("%H:%M:%S", time.gmtime(elapsed_time))))
# Close logger file
close_logger(logger)
def main(args):
train_loader = DataLoader(dataset=TrainDataset(),
batch_size=args.batch_size,
shuffle=True)
val_loader = DataLoader(dataset=ValDataset(), batch_size=args.batch_size)
model = LPRModel()
model = model.to(device)
loss_fn = torch.nn.CrossEntropyLoss()
optim = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
train_writer = SummaryWriter(logdir=os.path.join(args.log_dir, 'train'))
val_writer = SummaryWriter(os.path.join(args.log_dir, 'val'))
for epoch in range(1, args.epochs + 1):
avg_loss = 0
avg_acc = 0
model.train()
with tqdm(train_loader, ncols=0, leave=False) as pbar:
pbar.set_description('Epoch %3d/%3d' % (epoch, args.epochs))
for image, label in pbar:
image = image.to(device)
label = label.to(device)
predict = model(image)
loss = loss_fn(predict.transpose(1, 2), label)
optim.zero_grad()
loss.backward()
optim.step()
predict = torch.argmax(predict, dim=-1)
acc = (label == predict).float().mean(dim=1).sum().cpu()
avg_loss += loss.item()
avg_acc += acc.item()
pbar.set_postfix(loss='%.4f' % loss)
avg_loss = avg_loss / len(train_loader)
avg_acc = avg_acc / len(train_loader.dataset)
train_writer.add_scalar('loss', avg_loss, epoch)
train_writer.add_scalar('acc', avg_acc, epoch)
avg_val_loss = 0
avg_val_acc = 0
model.eval()
with torch.no_grad():
for image, label in val_loader:
image = image.to(device)
label = label.to(device)
predict = model(image)
loss = loss_fn(predict.transpose(1, 2), label)
predict = torch.argmax(predict, dim=-1)
acc = (label == predict).float().mean(dim=1).sum().cpu()
avg_val_loss += loss.item()
avg_val_acc += acc.item()
avg_val_loss = avg_val_loss / len(val_loader)
avg_val_acc = avg_val_acc / len(val_loader.dataset)
val_writer.add_scalar('loss', avg_val_loss, epoch)
val_writer.add_scalar('acc', avg_val_acc, epoch)
print(
'Epoch %3d/%3d train_loss: %.4f, val_loss: %.4f, '
'train_acc: %.4f, val_acc: %.4f' %
(epoch, args.epochs, avg_loss, avg_val_loss, avg_acc, avg_val_acc))
if epoch % args.checkpoint_period == 0:
path = os.path.join(args.checkpoint_dir, 'ckpt%d.pt' % epoch)
os.makedirs(args.checkpoint_dir, exist_ok=True)
torch.save(model.state_dict(), path)
def main(args):
# import network architecture
builder = ModelBuilder()
model = builder.build_net(arch=args.id,
num_input=args.num_input,
num_classes=args.num_classes,
num_branches=args.num_branches,
padding_list=args.padding_list,
dilation_list=args.dilation_list)
device_ids = [0, 2]
model = torch.nn.DataParallel(model, device_ids=device_ids).cuda()
# model = torch.nn.DataParallel(model, device_ids=list(range(args.num_gpus))).cuda()
cudnn.benchmark = True
if args.resume:
if os.path.isfile(args.resume):
print("=> Loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
state_dict = checkpoint['state_dict']
model.load_state_dict(state_dict)
print("=> Loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
else:
raise Exception("=> No checkpoint found at '{}'".format(
args.resume))
# initialization
num_ignore = 0
margin = [args.crop_size[k] - args.center_size[k] for k in range(3)]
num_images = int(len(test_dir) / args.num_input)
dice_score = np.zeros([num_images, 3]).astype(float)
for i in range(num_images):
# load the images, label and mask
im = []
for j in range(args.num_input):
direct, _ = test_dir[args.num_input * i + j].split("\n")
name = direct
if j < args.num_input - 1:
image = nib.load(args.root_path + direct + '.gz').get_data()
image = np.expand_dims(image, axis=0)
im.append(image)
if j == 0:
mask = nib.load(args.root_path + direct +
"/mask.nii.gz").get_data()
else:
labels = nib.load(args.root_path + direct + '.gz').get_data()
images = np.concatenate(im, axis=0).astype(float)
# divide the input images input small image segments
# return the padding input images which can be divided exactly
image_pad, mask_pad, label_pad, num_segments, padding_index, index = segment(
images, mask, labels, args)
# initialize prediction for the whole image as background
labels_shape = list(labels.shape)
labels_shape.append(args.num_classes)
pred = np.zeros(labels_shape)
pred[:, :, :, 0] = 1
# initialize the prediction for a small segmentation as background
pad_shape = [
int(num_segments[k] * args.center_size[k]) for k in range(3)
]
pad_shape.append(args.num_classes)
pred_pad = np.zeros(pad_shape)
pred_pad[:, :, :, 0] = 1
# score_per_image stores the sum of each image
score_per_image = np.zeros([3, 3])
# iterate over the z dimension
for idz in range(num_segments[2]):
tf = ValDataset(image_pad, label_pad, mask_pad, num_segments, idz,
args)
test_loader = DataLoader(tf,
batch_size=args.batch_size,
shuffle=args.shuffle,
num_workers=args.num_workers,
pin_memory=False)
score_seg, pred_seg = test(test_loader, model, num_segments, args)
pred_pad[:, :, idz * args.center_size[2]:(idz + 1) *
args.center_size[2], :] = pred_seg
score_per_image += score_seg
# decide the start and end point in the original image
for k in range(3):
if index[0][k] == 0:
index[0][k] = int(margin[k] / 2 - padding_index[0][k])
else:
index[0][k] = int(margin[k] / 2 + index[0][k])
index[1][k] = int(
min(index[0][k] + num_segments[k] * args.center_size[k],
labels.shape[k]))
dist = [index[1][k] - index[0][k] for k in range(3)]
pred[index[0][0]:index[1][0], index[0][1]:index[1][1],
index[0][2]:index[1][2]] = pred_pad[:dist[0], :dist[1], :dist[2]]
if np.sum(score_per_image[0, :]) == 0 or np.sum(
score_per_image[1, :]) == 0 or np.sum(
score_per_image[2, :]) == 0:
num_ignore += 1
continue
# compute the Enhance, Core and Whole dice score
dice_score_per = [
2 * np.sum(score_per_image[k, 2]) /
(np.sum(score_per_image[k, 0]) + np.sum(score_per_image[k, 1]))
for k in range(3)
]
print(
'Image: %d, Enhance score: %.4f, Core score: %.4f, Whole score: %.4f'
% (i, dice_score_per[0], dice_score_per[1], dice_score_per[2]))
dice_score[i, :] = dice_score_per
if args.visualize:
vis = np.argmax(pred, axis=3)
vis = np.swapaxes(vis, 0, 2).astype(dtype=np.uint8)
visualize_result(name, vis, args)
count_image = num_images - num_ignore
dice_score = dice_score[:count_image, :]
mean_dice = np.mean(dice_score, axis=0)
std_dice = np.std(dice_score, axis=0)
print('Evalution Done!')
print(
'Enhance score: %.4f, Core score: %.4f, Whole score: %.4f, Mean Dice score: %.4f'
% (mean_dice[0], mean_dice[1], mean_dice[2], np.mean(mean_dice)))
print(
'Enhance std: %.4f, Core std: %.4f, Whole std: %.4f, Mean Std: %.4f' %
(std_dice[0], std_dice[1], std_dice[2], np.mean(std_dice)))
from torch.utils import data as Data
import torchvision
from torchsummary import summary
from modules import Siamese, Siamese_Ultra
from dataset import MyDataset, ValDataset, FewDataset
from torch.utils.data import DataLoader
from tqdm import tqdm
if __name__ == '__main__':
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
train_data = MyDataset()
#train_data = FewDataset(k=10,seed=1127)
val_data = ValDataset()
train_loader = DataLoader(train_data,
batch_size=200,
shuffle=True,
num_workers=0)
val_loader = DataLoader(val_data,
batch_size=1000,
shuffle=False,
num_workers=0)
EPOCH = 2000
PATIENCE = 100
BATCH_SIZE = 150
# ===== Loading pre_trained model =====
def main(**args):
r"""Performs the main training loop
"""
# Load dataset
print('> Loading datasets ...')
dataset_val = ValDataset(valsetdir=args['valset_dir'],
gtvalsetdir=args['gt_valset_dir'],
gray_mode=False)
loader_train = train_dali_loader(batch_size=args['batch_size'],\
file_root=args['trainset_dir'],\
gt_file_root=args['gt_trainset_dir'],\
sequence_length=args['temp_patch_size'],\
crop_size=args['patch_size'],\
epoch_size=args['max_number_patches'],\
device_id=args['device_id'][0],\
random_shuffle=True,\
temp_stride=3)
num_minibatches = int(args['max_number_patches'] // args['batch_size'])
ctrl_fr_idx = (args['temp_patch_size'] - 1) // 2
print("\t# of training samples: %d\n" % int(args['max_number_patches']))
# Init loggers
writer, logger = init_logging(args)
# Define GPU devices
device_ids = args['device_id']
torch.backends.cudnn.benchmark = True # CUDNN optimization
# Create model
model = FastDVDnet()
model = nn.DataParallel(model, device_ids=device_ids).cuda()
# Define loss
criterion = nn.MSELoss(reduction='sum')
criterion.cuda()
# Optimizer
optimizer = optim.Adam(model.parameters(), lr=args['lr'])
# Resume training or start anew
start_epoch, training_params = resume_training(args, model, optimizer)
# Training
start_time = time.time()
for epoch in range(start_epoch, args['epochs']):
# Set learning rate
current_lr, reset_orthog = lr_scheduler(epoch, args)
if reset_orthog:
training_params['no_orthog'] = True
# set learning rate in optimizer
for param_group in optimizer.param_groups:
param_group["lr"] = current_lr
print('\nlearning rate %f' % current_lr)
# train
for i, data in enumerate(loader_train, 0):
# Pre-training step
model.train()
# When optimizer = optim.Optimizer(net.parameters()) we only zero the optim's grads
optimizer.zero_grad()
# convert inp to [N, num_frames*C. H, W] in [0., 1.] from [N, num_frames, C. H, W] in [0., 255.]
# extract ground truth (central frame)
img_train, gt_train = normalize_augment(data[0], ctrl_fr_idx)
N, _, H, W = img_train.size()
# Send tensors to GPU
gt_train = gt_train.cuda(non_blocking=True)
img_train = img_train.cuda(non_blocking=True)
# Evaluate model and optimize it
out_train = model(img_train)
'''
while torch.isinf(out_train).any():
print("out_inf")
return
'''
loss = criterion(gt_train, out_train) / (N * 2)
loss.backward()
optimizer.step()
# Results
if training_params['step'] % args['save_every'] == 0:
# Apply regularization by orthogonalizing filters
if not training_params['no_orthog']:
model.apply(svd_orthogonalization)
# Compute training PSNR
log_train_psnr(out_train, \
gt_train, \
loss, \
writer, \
epoch, \
i, \
num_minibatches, \
training_params)
# update step counter
training_params['step'] += 1
# Call to model.eval() to correctly set the BN layers before inference
model.eval()
# Validation and log images
validate_and_log(
model_temp=model, \
dataset_val=dataset_val, \
valnoisestd=0, \
temp_psz=args['temp_patch_size'], \
writer=writer, \
epoch=epoch, \
lr=current_lr, \
logger=logger, \
trainimg=img_train
)
# save model and checkpoint
training_params['start_epoch'] = epoch + 1
save_model_checkpoint(model, args, optimizer, training_params, epoch)
# Print elapsed time
elapsed_time = time.time() - start_time
print('Elapsed time {}'.format(
time.strftime("%H:%M:%S", time.gmtime(elapsed_time))))
# Close logger file
close_logger(logger)
def main(cfg, gpus):
torch.backends.cudnn.enabled = False
# cudnn.deterministic = False
# cudnn.enabled = True
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = ModelBuilder.build_decoder(
arch=cfg.MODEL.arch_decoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
num_class=cfg.DATASET.num_class,
weights=cfg.MODEL.weights_decoder)
if cfg.MODEL.arch_decoder == 'ocr':
print('Using cross entropy loss')
crit = CrossEntropy(ignore_label=-1)
else:
crit = nn.NLLLoss(ignore_index=-1)
if cfg.MODEL.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit,
cfg.TRAIN.deep_sup_scale)
else:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit)
# Dataset and Loader
dataset_train = TrainDataset(cfg.DATASET.root_dataset,
cfg.DATASET.list_train,
cfg.DATASET,
batch_per_gpu=cfg.TRAIN.batch_size_per_gpu)
loader_train = torch.utils.data.DataLoader(
dataset_train,
batch_size=len(gpus),
shuffle=False, # parameter is not used
collate_fn=user_scattered_collate,
num_workers=cfg.TRAIN.workers,
drop_last=True,
pin_memory=True)
# create loader iterator
iterator_train = iter(loader_train)
print('1 Epoch = {} iters'.format(cfg.TRAIN.epoch_iters))
if cfg.TRAIN.eval:
# Dataset and Loader for validtaion data
dataset_val = ValDataset(cfg.DATASET.root_dataset,
cfg.DATASET.list_val, cfg.DATASET)
loader_val = torch.utils.data.DataLoader(
dataset_val,
batch_size=cfg.VAL.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
iterator_val = iter(loader_val)
# load nets into gpu
if len(gpus) > 1:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=gpus)
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder, crit)
optimizers = create_optimizers(nets, cfg)
# Main loop
history = {
'train': {
'epoch': [],
'loss': [],
'acc': [],
'last_score': 0,
'best_score': cfg.TRAIN.best_score
}
}
for epoch in range(cfg.TRAIN.start_epoch, cfg.TRAIN.num_epoch):
train(segmentation_module, iterator_train, optimizers, history,
epoch + 1, cfg)
# calculate segmentation score
if cfg.TRAIN.eval and epoch in range(cfg.TRAIN.start_epoch,
cfg.TRAIN.num_epoch,
step=cfg.TRAIN.eval_step):
iou, acc = evaluate(segmentation_module, iterator_val, cfg, gpus)
history['train']['last_score'] = (iou + acc) / 2
if history['train']['last_score'] > history['train']['best_score']:
history['train']['best_score'] = history['train']['last_score']
checkpoint(nets, history, cfg, 'best_score')
# checkpointing
checkpoint(nets, history, cfg, epoch + 1)
print('Training Done!')
def main(cfg, gpu):
"""Adversarial examples for semantic segmentation models evaluated over all images in a list file."""
# if check_all_done(cfg):
# print("Processing", cfg.ATTACK.list_attack)
# print("Arguments", sys.argv)
# print("Entire experiment is already done. Quitting")
# return
# Save the command line args used to run the program
f_cmdline = open(
os.path.join(
cfg.ATTACK.output_dir, 'cmdline' +
str(datetime.datetime.now()).replace(' ', '_') + '.txt'), 'w')
for i in range(0, len(sys.argv)):
f_cmdline.write(sys.argv[i] + " ")
f_cmdline.close()
torch.cuda.device(gpu)
cfg.MODEL.model_name = os.path.basename(cfg.MODEL.arch_encoder + '_' +
cfg.MODEL.arch_decoder)
cfg.MODEL.arch_encoder = cfg.MODEL.arch_encoder.lower()
cfg.MODEL.arch_decoder = cfg.MODEL.arch_decoder.lower()
net_encoder = ModelBuilder.build_encoder(arch=cfg.MODEL.arch_encoder,
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = ModelBuilder.build_decoder(arch=cfg.MODEL.arch_decoder,
fc_dim=cfg.MODEL.fc_dim,
num_class=cfg.DATASET.num_class,
weights=cfg.MODEL.weights_decoder,
use_softmax=True)
# Copy the model prototxt to the folder
# shutil.copyfile(args.model_def, os.path.join(args.out_dir, model_name + '.prototxt'))
# Create the network and start the actual experiment
criterion = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder,
criterion)
# image_names = open(args.image_file, 'r').readlines()
# image_names = [x.strip() for x in image_names]
#
# for i, im_name in enumerate(image_names):
# args.image = im_name
# main_seg(args, net)
#
# if i % args.iter_print == 0:
# time_str = str(datetime.datetime.now())
# print("[{}] Image {}: {}".format(time_str, i, im_name))
# sys.stdout.flush()
# Dataset and Loader
orig_dataset = ValDataset(cfg.DATASET.root_dataset, cfg.ATTACK.list_attack,
cfg.DATASET)
orig_loader = torch.utils.data.DataLoader(
orig_dataset,
batch_size=cfg.ATTACK.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
# image_name = os.path.basename(args.image).split('.')[0].replace('_leftImg8bit', '')
# TODO: Check if required and modify accordingly
# if check_already_processed(cfg, image_name, model_name):
# return
# if net is None:
# net = caffe.Net(args.model_def, args.model_weights, caffe.TEST)
# image, im_height, im_width, orig_image = lib.PreprocessImage(
# args.image, args.pad_size, pad_value=args.pad_value, resize_dims=args.resize_dims, args=args)
# orig_pred, orig_conf = PredictWrapper(
# net, image, orig_image, dummy_label=None, is_seg=True, args=args)
orig_pred, _, orig_results = evaluate(segmentation_module, orig_loader,
cfg, 'orig', gpu)
attack_dataset = AttackDataset(cfg.DATASET.root_dataset,
cfg.ATTACK.list_attack, cfg.DATASET)
adv_func_args = get_adv_func_args(cfg, segmentation_module, attack_dataset,
gpu)
list_adv, _ = adv_attacks[cfg.ATTACK.method](**adv_func_args)
pert_dataset = ValDataset(cfg.DATASET.root_dataset, list_adv, cfg.DATASET)
pert_loader = torch.utils.data.DataLoader(
pert_dataset,
batch_size=cfg.ATTACK.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
adv_pred, _, adv_results = evaluate(segmentation_module, pert_loader, cfg,
'pert', gpu)
print('Segmentation attack processing complete!')
def main(args):
torch.cuda.set_device(args.gpu)
# Network Builders
builder = ModelBuilder()
net_encoder = None
net_decoder = None
unet = None
if args.unet == False:
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
num_class=args.num_class,
weights=args.weights_decoder,
use_softmax=True)
else:
unet = builder.build_unet(num_class=args.num_class,
arch=args.arch_unet,
weights=args.weights_unet,
use_softmax=True)
crit = nn.NLLLoss()
if args.unet == False:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit)
else:
segmentation_module = SegmentationModule(net_encoder,
net_decoder,
crit,
is_unet=args.unet,
unet=unet)
'''
# Dataset and Loader
dataset_val = dl.loadVal()
loader_val = torchdata.DataLoader(
dataset_val,
batch_size=5,
shuffle=False,
num_workers=1,
drop_last=True)
'''
# Dataset and Loader
dataset_val = ValDataset(args.list_val, args, max_sample=args.num_val)
loader_val = torchdata.DataLoader(dataset_val,
batch_size=args.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
# Main loop
evaluate(segmentation_module, loader_val, args)
print('Evaluation Done!')
arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = ModelBuilder.build_decoder(
arch=cfg.MODEL.arch_decoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
num_class=cfg.DATASET.num_class,
weights=cfg.MODEL.weights_decoder,
use_softmax=True)
crit = nn.NLLLoss(ignore_index=-1)
segmentation_module = SegmentationModule(net_encoder, net_decoder, crit)
# Dataset and Loader
dataset_val = ValDataset(cfg.DATASET.root_dataset, cfg.DATASET.list_val,
cfg.DATASET)
loader_val = torch.utils.data.DataLoader(dataset_val,
batch_size=BATCH_SIZE,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
segmentation_module.cuda()
# Main loop
evaluate(segmentation_module, loader_val, cfg, gpu, model['model_name'],
model['arxiv_id'])
print('Evaluation Done!')
def main(cfg, gpu, eval=False, interactive=False):
global net_encoder, net_decoder, segmentation_module, \
dataset_train, loader_train, iterator_train, history, epoch, \
dataset_val, loader_val
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder).cuda()
net_decoder = BayesianC1_2(
ModelBuilder.build_decoder(arch=cfg.MODEL.arch_decoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
num_class=cfg.DATASET.num_class,
weights=cfg.MODEL.weights_decoder).cuda())
segmentation_module = BayesianSegmentationModule(net_encoder, net_decoder)
if eval or interactive:
load_checkpoint(net_decoder, cfg, 10)
# Dataset and Loader
if not eval or interactive:
dataset_train = TrainDataset(
cfg.DATASET.root_dataset,
cfg.DATASET.list_train,
cfg.DATASET,
batch_per_gpu=cfg.TRAIN.batch_size_per_gpu)
loader_train = torch.utils.data.DataLoader(
dataset_train,
batch_size=1,
shuffle=False,
collate_fn=lambda batch: batch[0],
num_workers=cfg.TRAIN.workers,
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(cfg.TRAIN.epoch_iters))
iterator_train = iter(loader_train)
if eval or interactive:
dataset_val = ValDataset(cfg.DATASET.root_dataset,
cfg.DATASET.list_val, cfg.DATASET)
loader_val = torch.utils.data.DataLoader(
dataset_val,
batch_size=cfg.VAL.batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
if interactive:
return
if eval:
evaluate(segmentation_module, loader_val, cfg, gpu)
print('Evaluation Done!')
else:
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
for epoch in range(cfg.BAYESIAN.start_epoch, cfg.BAYESIAN.num_epoch):
train(segmentation_module, iterator_train, history, epoch + 1, cfg,
gpu)
# checkpointing
checkpoint(net_decoder, history, cfg, epoch + 1)
print('Training Done!')