def main(args):
# Network Builders
builder = ModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
fc_dim=args.fc_dim,
weights=args.weights_encoder)
nr_classes = broden_dataset.nr.copy()
nr_classes['part'] = sum(
[len(parts) for obj, parts in broden_dataset.object_part.items()])
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=args.fc_dim,
nr_classes=nr_classes,
weights=args.weights_decoder)
# TODO(LYC):: move criterion outside model.
# crit = nn.NLLLoss(ignore_index=-1)
if args.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(net_encoder, net_decoder,
args.deep_sup_scale)
else:
segmentation_module = SegmentationModule(net_encoder, net_decoder)
print('1 Epoch = {} iters'.format(args.epoch_iters))
# create loader iterator
iterator_train = create_multi_source_train_data_loader(args=args)
# load nets into gpu
if args.num_gpus > 1:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=range(
args.num_gpus))
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder)
optimizers = create_optimizers(nets, args)
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
for epoch in range(args.start_epoch, args.num_epoch + 1):
train(segmentation_module, iterator_train, optimizers, history, epoch,
args)
# checkpointing
checkpoint(nets, history, args, epoch)
print('Training Done!')
def main(args):
# 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)
# TODO(LYC):: move criterion outside model.
# crit = nn.NLLLoss(ignore_index=-1)
if args.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(
net_encoder, net_decoder, args.deep_sup_scale)
else:
segmentation_module = SegmentationModule(
net_encoder, net_decoder)
print('1 Epoch = {} iters'.format(args.epoch_iters))
# create loader iterator
iterator_train = create_multi_source_train_data_loader(args=args)
# load nets into gpu
if args.num_gpus > 1:
segmentation_module = UserScatteredDataParallel(
segmentation_module,
device_ids=range(args.num_gpus))
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder)
optimizers = create_optimizers(nets, args)
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
for epoch in range(args.start_epoch, args.num_epoch + 1):
train(segmentation_module, iterator_train, optimizers, history, epoch, args)
# checkpointing
checkpoint(nets, history, args, epoch)
print('Training Done!')
def main(args):
# Network Builders
builder = ModelBuilder()
unet = builder.build_unet(num_class=args.num_class,
arch=args.unet_arch,
weights=args.weights_unet)
print("Froze the following layers: ")
for name, p in unet.named_parameters():
if p.requires_grad == False:
print(name)
print()
crit = DualLoss(mode="train")
segmentation_module = SegmentationModule(crit, unet)
train_augs = Compose([PaddingCenterCrop(256), RandomHorizontallyFlip(), RandomVerticallyFlip(), RandomRotate(180)])
test_augs = Compose([PaddingCenterCrop(256)])
# Dataset and Loader
dataset_train = AC17( #Loads 3D volumes
root=args.data_root,
split='train',
k_split=args.k_split,
augmentations=train_augs,
img_norm=args.img_norm)
ac17_train = load2D(dataset_train, split='train', deform=True) #Dataloader for 2D slices. Requires 3D loader.
loader_train = data.DataLoader(
ac17_train,
batch_size=args.batch_size_per_gpu,
shuffle=True,
num_workers=int(args.workers),
drop_last=True,
pin_memory=True)
dataset_val = AC17(
root=args.data_root,
split='val',
k_split=args.k_split,
augmentations=test_augs,
img_norm=args.img_norm)
ac17_val = load2D(dataset_val, split='val', deform=False)
loader_val = data.DataLoader(
ac17_val,
batch_size=1,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
# load nets into gpu
if len(args.gpus) > 1:
segmentation_module = UserScatteredDataParallel(
segmentation_module,
device_ids=args.gpus)
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder, crit) if args.unet == False else (unet, crit)
optimizers = create_optimizers(nets, args)
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': [], 'jaccard': []}}
best_val = {'epoch_1': 0, 'mIoU_1': 0,
'epoch_2': 0, 'mIoU_2': 0,
'epoch_3': 0, 'mIoU_3': 0,
'epoch' : 0, 'mIoU': 0}
for epoch in range(args.start_epoch, args.num_epoch + 1):
train(segmentation_module, loader_train, optimizers, history, epoch, args)
iou, loss = eval(loader_val, segmentation_module, args, crit)
#checkpointing
ckpted = False
if loss < 0.215:
ckpted = True
if iou[0] > best_val['mIoU_1']:
best_val['epoch_1'] = epoch
best_val['mIoU_1'] = iou[0]
ckpted = True
if iou[1] > best_val['mIoU_2']:
best_val['epoch_2'] = epoch
best_val['mIoU_2'] = iou[1]
ckpted = True
if iou[2] > best_val['mIoU_3']:
best_val['epoch_3'] = epoch
best_val['mIoU_3'] = iou[2]
ckpted = True
if (iou[0]+iou[1]+iou[2])/3 > best_val['mIoU']:
best_val['epoch'] = epoch
best_val['mIoU'] = (iou[0]+iou[1]+iou[2])/3
ckpted = True
if epoch % 50 == 0:
checkpoint(nets, history, args, epoch)
continue
if epoch == args.num_epoch:
checkpoint(nets, history, args, epoch)
continue
if epoch < 15:
ckpted = False
if ckpted == False:
continue
else:
checkpoint(nets, history, args, epoch)
continue
print()
print('Training Done!')
def main(args):
# 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)
else:
unet = builder.build_unet(num_class=args.num_class,
arch=args.unet_arch,
weights=args.weights_unet)
print("Froze the following layers: ")
for name, p in unet.named_parameters():
if p.requires_grad == False:
print(name)
print()
crit = ACLoss(mode="train")
#crit = nn.CrossEntropyLoss().cuda()
#crit = nn.BCEWithLogitsLoss(pos_weight=torch.tensor(50))
#crit = nn.CrossEntropyLoss().cuda()
#crit = nn.BCELoss()
if args.arch_decoder.endswith('deepsup') and args.unet == False:
segmentation_module = SegmentationModule(
net_encoder, net_decoder, crit, args.deep_sup_scale)
else:
segmentation_module = SegmentationModule(
net_encoder, net_decoder, crit, is_unet=args.unet, unet=unet)
train_augs = Compose([PaddingCenterCrop(256), RandomHorizontallyFlip(), RandomVerticallyFlip(), RandomRotate(180)])
test_augs = Compose([PaddingCenterCrop(256)])
# Dataset and Loader
dataset_train = AC17(
root=args.data_root,
split='train',
k_split=args.k_split,
augmentations=train_augs,
img_norm=args.img_norm)
ac17_train = load2D(dataset_train, split='train', deform=True)
loader_train = data.DataLoader(
ac17_train,
batch_size=args.batch_size_per_gpu, # we have modified data_parallel
shuffle=True,
num_workers=int(args.workers),
drop_last=True,
pin_memory=True)
dataset_val = AC17(
root=args.data_root,
split='val',
k_split=args.k_split,
augmentations=test_augs,
img_norm=args.img_norm)
ac17_val = load2D(dataset_val, split='val', deform=False)
loader_val = data.DataLoader(
ac17_val,
batch_size=1,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
# create loader iterator
#iterator_train = iter(loader_train)
# load nets into gpu
if len(args.gpus) > 1:
segmentation_module = UserScatteredDataParallel(
segmentation_module,
device_ids=args.gpus)
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder, crit) if args.unet == False else (unet, crit)
optimizers = create_optimizers(nets, args)
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': [], 'jaccard': []}}
best_val = {'epoch_1': 0, 'mIoU_1': 0,
'epoch_2': 0, 'mIoU_2': 0,
'epoch_3': 0, 'mIoU_3': 0,
'epoch' : 0, 'mIoU': 0}
for epoch in range(args.start_epoch, args.num_epoch + 1):
train(segmentation_module, loader_train, optimizers, history, epoch, args)
iou, loss = eval(loader_val, segmentation_module, args, crit)
#checkpointing
ckpted = False
if loss < 0.215:
ckpted = True
if iou[0] > best_val['mIoU_1']:
best_val['epoch_1'] = epoch
best_val['mIoU_1'] = iou[0]
ckpted = True
if iou[1] > best_val['mIoU_2']:
best_val['epoch_2'] = epoch
best_val['mIoU_2'] = iou[1]
ckpted = True
if iou[2] > best_val['mIoU_3']:
best_val['epoch_3'] = epoch
best_val['mIoU_3'] = iou[2]
ckpted = True
if (iou[0]+iou[1]+iou[2])/3 > best_val['mIoU']:
best_val['epoch'] = epoch
best_val['mIoU'] = (iou[0]+iou[1]+iou[2])/3
ckpted = True
if epoch % 50 == 0:
checkpoint(nets, history, args, epoch)
continue
if epoch == args.num_epoch:
checkpoint(nets, history, args, epoch)
continue
if epoch < 15:
ckpted = False
if ckpted == False:
continue
else:
checkpoint(nets, history, args, epoch)
continue
print()
#print("[Val] Class 1: Epoch " + str(best_val['epoch_1']) + " had the best mIoU of " + str(best_val['mIoU_1']) + ".")
#print("[Val] Class 2: Epoch " + str(best_val['epoch_2']) + " had the best mIoU of " + str(best_val['mIoU_2']) + ".")
#print("[Val] Class 3: Epoch " + str(best_val['epoch_3']) + " had the best mIoU of " + str(best_val['mIoU_3']) + ".")
print('Training Done!')
def main():
"""Create the model and start the training."""
with open(args.config) as f:
config = yaml.load(f)
for k, v in config['common'].items():
setattr(args, k, v)
mkdirs(osp.join("logs/"+args.exp_name))
logger = create_logger('global_logger', "logs/" + args.exp_name + '/log.txt')
logger.info('{}'.format(args))
##############################
for key, val in vars(args).items():
logger.info("{:16} {}".format(key, val))
logger.info("random_scale {}".format(args.random_scale))
logger.info("is_training {}".format(args.is_training))
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
h, w = map(int, args.input_size_target.split(','))
input_size_target = (h, w)
print(type(input_size_target[1]))
cudnn.enabled = True
args.snapshot_dir = args.snapshot_dir + args.exp_name
tb_logger = SummaryWriter("logs/"+args.exp_name)
##############################
#validation data
h, w = map(int, args.input_size_test.split(','))
input_size_test = (h,w)
h, w = map(int, args.com_size.split(','))
com_size = (h, w)
h, w = map(int, args.input_size_crop.split(','))
input_size_crop = h,w
h,w = map(int, args.input_size_target_crop.split(','))
input_size_target_crop = h,w
test_normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
test_transform = transforms.Compose([
transforms.Resize((input_size_test[1], input_size_test[0])),
transforms.ToTensor(),
test_normalize])
valloader = data.DataLoader(cityscapesDataSet(
args.data_dir_target,
args.data_list_target_val,
crop_size=input_size_test,
set='train',
transform=test_transform),num_workers=args.num_workers,
batch_size=1, shuffle=False, pin_memory=True)
with open('./dataset/cityscapes_list/info.json', 'r') as fp:
info = json.load(fp)
mapping = np.array(info['label2train'], dtype=np.int)
label_path_list_val = args.label_path_list_val
label_path_list_test = args.label_path_list_test
label_path_list_test = './dataset/cityscapes_list/label.txt'
gt_imgs_val = open(label_path_list_val, 'r').read().splitlines()
gt_imgs_val = [osp.join(args.data_dir_target_val, x) for x in gt_imgs_val]
testloader = data.DataLoader(cityscapesDataSet(
args.data_dir_target,
args.data_list_target_test,
crop_size=input_size_test,
set='val',
transform=test_transform),
num_workers=args.num_workers,
batch_size=1,
shuffle=False, pin_memory=True)
gt_imgs_test = open(label_path_list_test ,'r').read().splitlines()
gt_imgs_test = [osp.join(args.data_dir_target_test, x) for x in gt_imgs_test]
name_classes = np.array(info['label'], dtype=np.str)
interp_val = nn.Upsample(size=(com_size[1], com_size[0]),mode='bilinear', align_corners=True)
####
#build model
####
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_classes,
weights=args.weights_decoder,
use_aux=True)
model = SegmentationModule(
net_encoder, net_decoder, args.use_aux)
if args.num_gpus > 1:
model = torch.nn.DataParallel(model)
patch_replication_callback(model)
model.cuda()
nets = (net_encoder, net_decoder, None, None)
optimizers = create_optimizer(nets, args)
cudnn.enabled=True
cudnn.benchmark=True
model.train()
mean=[0.485, 0.456, 0.406]
std=[0.229, 0.224, 0.225]
source_normalize = transforms_seg.Normalize(mean=mean,
std=std)
mean_mapping = [0.485, 0.456, 0.406]
mean_mapping = [item * 255 for item in mean_mapping]
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
source_transform = transforms_seg.Compose([
transforms_seg.Resize([input_size[1], input_size[0]]),
segtransforms.RandScale((args.scale_min, args.scale_max)),
#segtransforms.RandRotate((args.rotate_min, args.rotate_max), padding=mean_mapping, ignore_label=args.ignore_label),
#segtransforms.RandomGaussianBlur(),
segtransforms.RandomHorizontalFlip(),
segtransforms.Crop([input_size_crop[1], input_size_crop[0]], crop_type='rand', padding=mean_mapping, ignore_label=args.ignore_label),
transforms_seg.ToTensor(),
source_normalize])
target_normalize = transforms_seg.Normalize(mean=mean,
std=std)
target_transform = transforms_seg.Compose([
transforms_seg.Resize([input_size_target[1], input_size_target[0]]),
segtransforms.RandScale((args.scale_min, args.scale_max)),
#segtransforms.RandRotate((args.rotate_min, args.rotate_max), padding=mean_mapping, ignore_label=args.ignore_label),
#segtransforms.RandomGaussianBlur(),
segtransforms.RandomHorizontalFlip(),
segtransforms.Crop([input_size_target_crop[1], input_size_target_crop[0]],crop_type='rand', padding=mean_mapping, ignore_label=args.ignore_label),
transforms_seg.ToTensor(),
target_normalize])
trainloader = data.DataLoader(
GTA5DataSet(args.data_dir, args.data_list, max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size, transform = source_transform),
batch_size=args.batch_size, shuffle=True, num_workers=1, pin_memory=True)
trainloader_iter = enumerate(trainloader)
targetloader = data.DataLoader(fake_cityscapesDataSet(args.data_dir_target, args.data_list_target,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
set=args.set,
transform=target_transform),
batch_size=args.batch_size, shuffle=True, num_workers=1,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
# implement model.optim_parameters(args) to handle different models' lr setting
criterion_seg = torch.nn.CrossEntropyLoss(ignore_index=255,reduce=False)
interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), align_corners=True, mode='bilinear')
# labels for adversarial training
source_label = 0
target_label = 1
optimizer_encoder, optimizer_decoder, optimizer_disc, optimizer_reconst = optimizers
batch_time = AverageMeter(10)
loss_seg_value1 = AverageMeter(10)
is_best_test = True
best_mIoUs = 0
loss_seg_value2 = AverageMeter(10)
loss_balance_value = AverageMeter(10)
loss_pseudo_value = AverageMeter(10)
bounding_num = AverageMeter(10)
pseudo_num = AverageMeter(10)
for i_iter in range(args.num_steps):
# train G
# don't accumulate grads in D
end = time.time()
_, batch = trainloader_iter.__next__()
images, labels, _ = batch
images = Variable(images).cuda(async=True)
labels = Variable(labels).cuda(async=True)
seg, aux_seg, loss_seg2, loss_seg1 = model(images, labels)
loss_seg2 = torch.mean(loss_seg2)
loss_seg1 = torch.mean(loss_seg1)
loss = loss_seg2+args.lambda_seg*loss_seg1
#logger.info(loss_seg1.data.cpu().numpy())
loss_seg_value2.update(loss_seg2.data.cpu().numpy())
# train with target
optimizer_encoder.zero_grad()
optimizer_decoder.zero_grad()
loss.backward()
optimizer_encoder.step()
optimizer_decoder.step()
del seg, loss_seg2
_, batch = targetloader_iter.__next__()
with torch.no_grad():
images, labels, _ = batch
images = Variable(images).cuda(async=True)
result = model(images, None)
del result
batch_time.update(time.time() - end)
remain_iter = args.num_steps - i_iter
remain_time = remain_iter * batch_time.avg
t_m, t_s = divmod(remain_time, 60)
t_h, t_m = divmod(t_m, 60)
remain_time = '{:02d}:{:02d}:{:02d}'.format(int(t_h), int(t_m), int(t_s))
adjust_learning_rate(optimizer_encoder, i_iter, args.lr_encoder, args)
adjust_learning_rate(optimizer_decoder, i_iter, args.lr_decoder, args)
if i_iter % args.print_freq == 0:
lr_encoder = optimizer_encoder.param_groups[0]['lr']
lr_decoder = optimizer_decoder.param_groups[0]['lr']
logger.info('exp = {}'.format(args.snapshot_dir))
logger.info('Iter = [{0}/{1}]\t'
'Time = {batch_time.avg:.3f}\t'
'loss_seg1 = {loss_seg1.avg:4f}\t'
'loss_seg2 = {loss_seg2.avg:.4f}\t'
'lr_encoder = {lr_encoder:.8f} lr_decoder = {lr_decoder:.8f}'.format(
i_iter, args.num_steps, batch_time=batch_time,
loss_seg1=loss_seg_value1, loss_seg2=loss_seg_value2,
lr_encoder=lr_encoder,
lr_decoder=lr_decoder))
logger.info("remain_time: {}".format(remain_time))
if not tb_logger is None:
tb_logger.add_scalar('loss_seg_value1', loss_seg_value1.avg, i_iter)
tb_logger.add_scalar('loss_seg_value2', loss_seg_value2.avg, i_iter)
tb_logger.add_scalar('lr', lr_encoder, i_iter)
#####
#save image result
if i_iter % args.save_pred_every == 0 and i_iter != 0:
logger.info('taking snapshot ...')
model.eval()
val_time = time.time()
hist = np.zeros((19,19))
for index, batch in tqdm(enumerate(valloader)):
with torch.no_grad():
image, name = batch
output2, _ = model(Variable(image).cuda(), None)
pred = interp_val(output2)
del output2
pred = pred.cpu().data[0].numpy()
pred = pred.transpose(1, 2, 0)
pred = np.asarray(np.argmax(pred, axis=2), dtype=np.uint8)
label = np.array(Image.open(gt_imgs_val[index]))
#label = np.array(label.resize(com_size, Image.
label = label_mapping(label, mapping)
#logger.info(label.shape)
hist += fast_hist(label.flatten(), pred.flatten(), 19)
mIoUs = per_class_iu(hist)
for ind_class in range(args.num_classes):
logger.info('===>' + name_classes[ind_class] + ':\t' + str(round(mIoUs[ind_class] * 100, 2)))
tb_logger.add_scalar(name_classes[ind_class] + '_mIoU', mIoUs[ind_class], i_iter)
mIoUs = round(np.nanmean(mIoUs) *100, 2)
if mIoUs >= best_mIoUs:
is_best_test = True
best_mIoUs = mIoUs
else:
is_best_test = False
logger.info("current mIoU {}".format(mIoUs))
logger.info("best mIoU {}".format(best_mIoUs))
tb_logger.add_scalar('val mIoU', mIoUs, i_iter)
tb_logger.add_scalar('val mIoU', mIoUs, i_iter)
net_encoder, net_decoder, net_disc, net_reconst = nets
save_checkpoint(net_encoder, 'encoder', i_iter, args, is_best_test)
save_checkpoint(net_decoder, 'decoder', i_iter, args, is_best_test)
model.train()
def main(cfg, gpus):
# 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_encoder = ModelBuilder.build_encoder(
# arch=cfg.MODEL.arch_encoder.lower(),
# fc_dim=cfg.MODEL.fc_dim,
# weights="pretrained/encoder_epoch_20.pth")
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)
# net_decoder = ModelBuilder.build_decoder(
# arch=cfg.MODEL.arch_decoder.lower(),
# fc_dim=cfg.MODEL.fc_dim,
# num_class=cfg.DATASET.num_class,
# weights="pretrained/decoder_epoch_20.pth")
ct = 0
# for child in net_encoder.children():
# ct+= 1
# if ct < 12:
# for parameter in net_encoder.layer1.parameters():
# parameter.requires_grad = False
# for parameter in net_encoder.layer2.parameters():
# parameter.requires_grad = False
# for parameter in net_encoder.parameters():
# parameter.requires_grad = False;
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), # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=cfg.TRAIN.workers,
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(cfg.TRAIN.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
# 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': []}}
for epoch in range(cfg.TRAIN.start_epoch, cfg.TRAIN.num_epoch):
train(segmentation_module, iterator_train, optimizers, history,
epoch + 1, cfg)
n = ((int(cfg.TRAIN.epoch_iters) // 20) * epoch) + 1
for i in zip(history['train']['loss'][n:],
history['train']['acc'][n:]):
writer.add_scalar('Loss/Train', i[0], epoch)
writer.add_scalar('Accuracy/Train', i[1], epoch)
# checkpointing
checkpoint(nets, history, cfg, epoch + 1)
writer.close()
print('Training Done!')
def main(cfg, gpus):
# Network Builders
net_encoder = ModelBuilder.build_encoder(
arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder,
dilate_rate=cfg.DATASET.segm_downsampling_rate)
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.foveation:
net_foveater = ModelBuilder.build_foveater(
in_channel=cfg.MODEL.in_dim,
out_channel=len(cfg.MODEL.patch_bank),
len_gpus=len(gpus),
weights=cfg.MODEL.weights_foveater,
cfg=cfg)
# tensor
writer = SummaryWriter('{}/tensorboard'.format(cfg.DIR))
if cfg.DATASET.root_dataset == '/scratch0/chenjin/GLEASON2019_DATA/Data/':
if cfg.TRAIN.loss_fun == 'DiceLoss':
crit = DiceLoss()
elif cfg.TRAIN.loss_fun == 'FocalLoss':
crit = FocalLoss()
elif cfg.TRAIN.loss_fun == 'DiceCoeff':
crit = DiceCoeff()
elif cfg.TRAIN.loss_fun == 'NLLLoss':
crit = nn.NLLLoss(ignore_index=-2)
else:
crit = OhemCrossEntropy(ignore_label=-1,
thres=0.9,
min_kept=100000,
weight=None)
elif 'ADE20K' in cfg.DATASET.root_dataset:
crit = nn.NLLLoss(ignore_index=-2)
elif 'CITYSCAPES' in cfg.DATASET.root_dataset:
if cfg.TRAIN.loss_fun == 'NLLLoss':
crit = nn.NLLLoss(ignore_index=19)
else:
class_weights = torch.FloatTensor([
0.8373, 0.918, 0.866, 1.0345, 1.0166, 0.9969, 0.9754, 1.0489,
0.8786, 1.0023, 0.9539, 0.9843, 1.1116, 0.9037, 1.0865, 1.0955,
1.0865, 1.1529, 1.0507
]).cuda()
crit = OhemCrossEntropy(ignore_label=20,
thres=0.9,
min_kept=131072,
weight=class_weights)
elif 'DeepGlob' in cfg.DATASET.root_dataset and (
cfg.TRAIN.loss_fun == 'FocalLoss'
or cfg.TRAIN.loss_fun == 'OhemCrossEntropy'):
if cfg.TRAIN.loss_fun == 'FocalLoss':
crit = FocalLoss(gamma=6, ignore_label=cfg.DATASET.ignore_index)
elif cfg.TRAIN.loss_fun == 'OhemCrossEntropy':
crit = OhemCrossEntropy(ignore_label=cfg.DATASET.ignore_index,
thres=0.9,
min_kept=131072)
else:
if cfg.TRAIN.loss_fun == 'NLLLoss':
if cfg.DATASET.ignore_index != -2:
crit = nn.NLLLoss(ignore_index=cfg.DATASET.ignore_index)
else:
crit = nn.NLLLoss(ignore_index=-2)
else:
if cfg.DATASET.ignore_index != -2:
crit = nn.CrossEntropyLoss(
ignore_index=cfg.DATASET.ignore_index)
else:
crit = nn.CrossEntropyLoss(ignore_index=-2)
# crit = DiceCoeff()
if cfg.MODEL.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit, cfg,
cfg.TRAIN.deep_sup_scale)
elif cfg.MODEL.foveation:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit, cfg)
else:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit, cfg)
if cfg.MODEL.foveation:
foveation_module = FovSegmentationModule(net_foveater,
cfg,
len_gpus=len(gpus))
total_fov = sum(
[param.nelement() for param in foveation_module.parameters()])
print('Number of FoveationModule params: %.2fM \n' % (total_fov / 1e6))
total = sum(
[param.nelement() for param in segmentation_module.parameters()])
print('Number of SegmentationModule params: %.2fM \n' % (total / 1e6))
# 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), # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=cfg.TRAIN.workers,
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(cfg.TRAIN.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
# load nets into gpu
if len(gpus) > 1:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=gpus)
# For sync bn
patch_replication_callback(segmentation_module)
if cfg.MODEL.foveation:
foveation_module = UserScatteredDataParallel(foveation_module,
device_ids=gpus)
patch_replication_callback(foveation_module)
segmentation_module.cuda()
if cfg.MODEL.foveation:
foveation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder, crit)
if cfg.MODEL.foveation:
nets = (net_encoder, net_decoder, crit, net_foveater)
optimizers = create_optimizers(nets, cfg)
# Main loop
if cfg.VAL.dice:
history = {
'train': {
'epoch': [],
'loss': [],
'acc': []
},
'save': {
'epoch': [],
'train_loss': [],
'train_acc': [],
'val_iou': [],
'val_dice': [],
'val_acc': [],
'print_grad': None
}
}
else:
history = {
'train': {
'epoch': [],
'loss': [],
'acc': []
},
'save': {
'epoch': [],
'train_loss': [],
'train_acc': [],
'val_iou': [],
'val_dice': [],
'val_acc': [],
'print_grad': None
}
}
if cfg.TRAIN.start_epoch > 0:
history_previous_epoches = pd.read_csv(
'{}/history_epoch_{}.csv'.format(cfg.DIR, cfg.TRAIN.start_epoch))
history['save']['epoch'] = list(history_previous_epoches['epoch'])
history['save']['train_loss'] = list(
history_previous_epoches['train_loss'])
history['save']['train_acc'] = list(
history_previous_epoches['train_acc'])
history['save']['val_iou'] = list(history_previous_epoches['val_iou'])
history['save']['val_acc'] = list(history_previous_epoches['val_acc'])
# if cfg.VAL.dice:
# history['save']['val_dice'] = history_previous_epoches['val_dice']
if not os.path.isdir(os.path.join(cfg.DIR,
"Fov_probability_distribution")):
os.makedirs(os.path.join(cfg.DIR, "Fov_probability_distribution"))
f_prob = []
for p in range(len(cfg.MODEL.patch_bank)):
f = open(
os.path.join(cfg.DIR, 'Fov_probability_distribution',
'patch_{}_distribution.txt'.format(p)), 'w')
f.close()
for epoch in range(cfg.TRAIN.start_epoch, cfg.TRAIN.num_epoch):
if cfg.MODEL.foveation:
train(segmentation_module, iterator_train, optimizers, epoch + 1,
cfg, history, foveation_module)
if history['train']['print_grad'] is not None and type(
history['train']['print_grad']) is not torch.Tensor:
if history['train']['print_grad'][
'layer1_grad'] is not None and history['train'][
'print_grad']['layer1_grad'][
history['train']['print_grad']['layer1_grad'] >
0].numel() > 0:
writer.add_histogram(
'Print non-zero gradient (layer1) histogram',
history['train']['print_grad']['layer1_grad'][
history['train']['print_grad']['layer1_grad'] > 0],
epoch + 1)
writer.add_histogram(
'Print gradient (layer1) histogram',
history['train']['print_grad']['layer1_grad'],
epoch + 1)
writer.add_scalar(
'Percentage none-zero gradients (layer1)',
history['train']['print_grad']['layer1_grad'][
history['train']['print_grad']['layer1_grad'] > 0].
numel() /
history['train']['print_grad']['layer1_grad'].numel(),
epoch + 1)
writer.add_image(
'Print_grad_Fov_softmax_layer1(normalized_b0_p0)',
(history['train']['print_grad']['layer1_grad'][0][0] -
history['train']['print_grad']['layer1_grad'][0]
[0].min()) /
(history['train']['print_grad']['layer1_grad'][0]
[0].max() - history['train']['print_grad']
['layer1_grad'][0][0].min()),
epoch + 1,
dataformats='HW')
if history['train']['print_grad'][
'layer2_grad'] is not None and history['train'][
'print_grad']['layer2_grad'][
history['train']['print_grad']['layer2_grad'] >
0].numel() > 0:
writer.add_histogram(
'Print non-zero gradient (layer2) histogram',
history['train']['print_grad']['layer2_grad'][
history['train']['print_grad']['layer2_grad'] > 0],
epoch + 1)
writer.add_histogram(
'Print gradient (layer2) histogram',
history['train']['print_grad']['layer2_grad'],
epoch + 1)
writer.add_scalar(
'Percentage none-zero gradients (layer2)',
history['train']['print_grad']['layer2_grad'][
history['train']['print_grad']['layer2_grad'] > 0].
numel() /
history['train']['print_grad']['layer2_grad'].numel(),
epoch + 1)
writer.add_image(
'Print_grad_Fov_softmax_layer2(normalized_b0_p0)',
(history['train']['print_grad']['layer2_grad'][0][0] -
history['train']['print_grad']['layer2_grad'][0]
[0].min()) /
(history['train']['print_grad']['layer2_grad'][0]
[0].max() - history['train']['print_grad']
['layer2_grad'][0][0].min()),
epoch + 1,
dataformats='HW')
if history['train']['print_grad'][
'layer3_grad'] is not None and history['train'][
'print_grad']['layer3_grad'][
history['train']['print_grad']['layer3_grad'] >
0].numel() > 0:
writer.add_histogram(
'Print non-zero gradient (layer3) histogram',
history['train']['print_grad']['layer3_grad'][
history['train']['print_grad']['layer3_grad'] > 0],
epoch + 1)
writer.add_histogram(
'Print gradient (layer3) histogram',
history['train']['print_grad']['layer3_grad'],
epoch + 1)
writer.add_scalar(
'Percentage none-zero gradients (layer3)',
history['train']['print_grad']['layer3_grad'][
history['train']['print_grad']['layer3_grad'] > 0].
numel() /
history['train']['print_grad']['layer3_grad'].numel(),
epoch + 1)
writer.add_image(
'Print_grad_Fov_softmax_layer3(normalized_b0_p0)',
(history['train']['print_grad']['layer3_grad'][0][0] -
history['train']['print_grad']['layer3_grad'][0]
[0].min()) /
(history['train']['print_grad']['layer3_grad'][0]
[0].max() - history['train']['print_grad']
['layer3_grad'][0][0].min()),
epoch + 1,
dataformats='HW')
else:
train(segmentation_module, iterator_train, optimizers, epoch + 1,
cfg, history)
# checkpointing
if (epoch + 1) % cfg.TRAIN.checkpoint_per_epoch == 0:
checkpoint(nets, cfg, epoch + 1)
checkpoint_last(nets, cfg, epoch + 1)
else:
checkpoint_last(nets, cfg, epoch + 1)
if (epoch + 1) % cfg.TRAIN.eval_per_epoch == 0:
# eval during train
if cfg.VAL.multipro:
if cfg.MODEL.foveation:
if cfg.VAL.all_F_Xlr_time:
val_iou, val_acc, F_Xlr_all, F_Xlr_score_flat_all = eval_during_train_multipro(
cfg, gpus)
else:
val_iou, val_acc, F_Xlr, F_Xlr_score_flat = eval_during_train_multipro(
cfg, gpus)
else:
val_iou, val_acc = eval_during_train_multipro(cfg, gpus)
else:
if cfg.VAL.dice:
if cfg.MODEL.foveation:
if cfg.VAL.all_F_Xlr_time:
val_iou, val_dice, val_acc, F_Xlr_all, F_Xlr_score_flat_all = eval_during_train(
cfg)
else:
val_iou, val_dice, val_acc, F_Xlr, F_Xlr_score_flat = eval_during_train(
cfg)
else:
val_iou, val_dice, val_acc = eval_during_train(cfg)
else:
if cfg.MODEL.foveation:
if cfg.VAL.all_F_Xlr_time:
val_iou, val_acc, F_Xlr_all, F_Xlr_score_flat_all = eval_during_train(
cfg)
else:
val_iou, val_acc, F_Xlr, F_Xlr_score_flat = eval_during_train(
cfg)
else:
val_iou, val_acc = eval_during_train(cfg)
history['save']['epoch'].append(epoch + 1)
history['save']['train_loss'].append(history['train']['loss'][-1])
history['save']['train_acc'].append(history['train']['acc'][-1] *
100)
history['save']['val_iou'].append(val_iou)
if cfg.VAL.dice:
history['save']['val_dice'].append(val_dice)
history['save']['val_acc'].append(val_acc)
# write to tensorboard
writer.add_scalar('Loss/train', history['train']['loss'][-1],
epoch + 1)
writer.add_scalar('Acc/train', history['train']['acc'][-1] * 100,
epoch + 1)
writer.add_scalar('Acc/val', val_acc, epoch + 1)
writer.add_scalar('mIoU/val', val_iou, epoch + 1)
if cfg.VAL.dice:
writer.add_scalar('mDice/val', val_acc, epoch + 1)
if cfg.VAL.all_F_Xlr_time:
print('=============F_Xlr_score_flat_all================\n',
F_Xlr_score_flat_all.shape)
for p in range(F_Xlr_score_flat_all.shape[0]):
# add small artifact to modify range, because no range flag in add_histogram
F_Xlr_score_flat_all[p][0] = 0
F_Xlr_score_flat_all[p][-1] = 1
writer.add_histogram(
'Patch_{} probability histogram'.format(p),
F_Xlr_score_flat_all[p], epoch + 1)
f = open(
os.path.join(cfg.DIR, 'Fov_probability_distribution',
'patch_{}_distribution.txt'.format(p)),
'a')
if epoch == 0:
f.write('epoch/ bins: {}\n'.format(
np.histogram(F_Xlr_score_flat_all[p],
bins=10,
range=(0, 1))[1]))
f.write('epoch {}: {}\n'.format(
epoch + 1,
np.histogram(F_Xlr_score_flat_all[p],
bins=10,
range=(0, 1))[0] / sum(
np.histogram(F_Xlr_score_flat_all[p],
bins=10,
range=(0, 1))[0])))
f.close()
writer.add_histogram('Patch_All probability histogram',
F_Xlr_score_flat_all, epoch + 1)
else:
for p in range(F_Xlr_score_flat_all.shape[0]):
F_Xlr_score_flat[p][0] = 0
F_Xlr_score_flat[p][-1] = 1
writer.add_histogram(
'Patch_{} probability histogram'.format(p),
F_Xlr_score_flat[p], epoch + 1)
writer.add_histogram('Patch_All probability histogram',
F_Xlr_score_flat, epoch + 1)
else:
history['save']['epoch'].append(epoch + 1)
history['save']['train_loss'].append(history['train']['loss'][-1])
history['save']['train_acc'].append(history['train']['acc'][-1] *
100)
history['save']['val_iou'].append('')
if cfg.VAL.dice:
history['save']['val_dice'].append('')
history['save']['val_acc'].append('')
# write to tensorboard
writer.add_scalar('Loss/train', history['train']['loss'][-1],
epoch + 1)
writer.add_scalar('Acc/train', history['train']['acc'][-1] * 100,
epoch + 1)
# writer.add_scalar('Acc/val', '', epoch+1)
# writer.add_scalar('mIoU/val', '', epoch+1)
# saving history
checkpoint_history(history, cfg, epoch + 1)
if (epoch + 1) % cfg.TRAIN.eval_per_epoch == 0:
# output F_Xlr
if cfg.MODEL.foveation:
# save time series F_Xlr (t,b,d,w,h)
if epoch == 0 or epoch == cfg.TRAIN.start_epoch:
if cfg.VAL.all_F_Xlr_time:
F_Xlr_time_all = []
for val_idx in range(len(F_Xlr_all)):
F_Xlr_time_all.append(F_Xlr_all[val_idx][0])
else:
F_Xlr_time = F_Xlr
else:
if cfg.VAL.all_F_Xlr_time:
for val_idx in range(len(F_Xlr_all)):
F_Xlr_time_all[val_idx] = np.concatenate(
(F_Xlr_time_all[val_idx],
F_Xlr_all[val_idx][0]),
axis=0)
else:
F_Xlr_time = np.concatenate((F_Xlr_time, F_Xlr),
axis=0)
if cfg.VAL.all_F_Xlr_time:
for val_idx in range(len(F_Xlr_all)):
print('F_Xlr_time_{}'.format(F_Xlr_all[val_idx][1]),
F_Xlr_time_all[val_idx].shape)
if not os.path.isdir(
os.path.join(cfg.DIR, "F_Xlr_time_all_vals")):
os.makedirs(
os.path.join(cfg.DIR, "F_Xlr_time_all_vals"))
np.save(
'{}/F_Xlr_time_all_vals/F_Xlr_time_last_{}.npy'.
format(cfg.DIR, F_Xlr_all[val_idx][1]),
F_Xlr_time_all[val_idx])
else:
print('F_Xlr_time', F_Xlr_time.shape)
np.save('{}/F_Xlr_time_last.npy'.format(cfg.DIR),
F_Xlr_time)
if not cfg.TRAIN.save_checkpoint:
os.remove('{}/encoder_epoch_last.pth'.format(cfg.DIR))
os.remove('{}/decoder_epoch_last.pth'.format(cfg.DIR))
print('Training Done!')
writer.close()
def main(cfg, gpus):
# Network Builders
net_enc_query = ModelBuilder.build_encoder(
arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_enc_query)
if cfg.MODEL.memory_encoder_noBN:
net_enc_memory = ModelBuilder.build_encoder_memory_separate(
arch=cfg.MODEL.arch_encoder.lower()+'_nobn',
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_enc_memory,
num_class=cfg.DATASET.num_class)
else:
net_enc_memory = ModelBuilder.build_encoder_memory_separate(
arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_enc_memory,
num_class=cfg.DATASET.num_class,
pretrained=cfg.memory_enc_pretrained)
net_att_query = ModelBuilder.build_encoder(
arch='attention',
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_att_query)
net_att_memory = ModelBuilder.build_encoder(
arch='attention',
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_att_memory)
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)
crit = nn.NLLLoss(ignore_index=-1)
if cfg.MODEL.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationAttentionSeparateModule(
net_enc_query, net_enc_memory, net_att_query, net_att_memory, net_decoder, crit, cfg.TRAIN.deep_sup_scale, zero_memory=cfg.MODEL.zero_memory, random_memory_bias=cfg.MODEL.random_memory_bias, random_memory_nobias=cfg.MODEL.random_memory_nobias, random_scale=cfg.MODEL.random_scale, qval_qread_BN=cfg.MODEL.qval_qread_BN, normalize_key=cfg.MODEL.normalize_key, p_scalar=cfg.MODEL.p_scalar, debug=True)
else:
segmentation_module = SegmentationAttentionSeparateModule(
net_enc_query, net_enc_memory, net_att_query, net_att_memory, net_decoder, crit, zero_memory=cfg.MODEL.zero_memory, random_memory_bias=cfg.MODEL.random_memory_bias, random_memory_nobias=cfg.MODEL.random_memory_nobias, random_scale=cfg.MODEL.random_scale, qval_qread_BN=cfg.MODEL.qval_qread_BN, normalize_key=cfg.MODEL.normalize_key, p_scalar=cfg.MODEL.p_scalar, debug=True)
# Dataset and Loader
dataset_train = TrainDataset(
cfg.DATASET.root_dataset,
cfg.DATASET.list_train,
cfg.DATASET,
cfg.DATASET.ref_path,
cfg.DATASET.ref_start,
cfg.DATASET.ref_end,
batch_per_gpu=cfg.TRAIN.batch_size_per_gpu)
loader_train = torch.utils.data.DataLoader(
dataset_train,
batch_size=len(gpus), # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=cfg.TRAIN.workers,
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(cfg.TRAIN.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
# 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 = UserScatteredDataParallel(
segmentation_module,
device_ids=gpus)
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_enc_query, net_enc_memory, net_att_query, net_att_memory, net_decoder, crit)
optimizers = create_optimizers(nets, cfg)
for epoch in range(0, 1):
train(segmentation_module, iterator_train, optimizers, None, epoch+1, cfg)
print('Training Done!')
for name, p in unet.named_parameters():
if p.requires_grad == False:
print(name)
print()
crit = DualLoss(mode="train")
segmentation_module = SegmentationModule(crit, unet)
# load nets into gpu
if len(args.gpus) > 1:
segmentation_module = UserScatteredDataParallel(
segmentation_module,
device_ids=args.gpus)
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
print("ready to load data")
import train
import streamlit as st
from os import listdir
from os.path import isfile, join
from PIL import Image
test_path = '/home/rexma/demo/LCTSC/demo/'
temp_path = '/home/rexma/demo/LCTSC/temp/'
result_path = '/home/rexma/demo/LCTSC/result/'
st.sidebar.title("About")
st.sidebar.info(
"This is a demo application written to help you understand Streamlit. The application identifies the animal in the picture. It was built using a Convolution Neural Network (CNN).")
def main(args):
# Network Builders
builder = VGGModelBuilder()
net_encoder = builder.build_encoder(arch=args.arch_encoder,
weights=args.weights_encoder)
net_decoder = builder.build_decoder(arch=args.arch_decoder,
fc_dim=1024,
num_class=args.num_class,
weights=args.weights_decoder)
crit = nn.CrossEntropyLoss(ignore_index=255, reduction='sum')
if args.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit, args.deep_sup_scale)
else:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit)
print(segmentation_module)
# Dataset and Loader
dataset_train = VOCTrainDataset(args,
batch_per_gpu=args.batch_size_per_gpu)
loader_train = torchdata.DataLoader(
dataset_train,
batch_size=1, # data_parallel have been modified, not useful
shuffle=False, # do not use this param
collate_fn=user_scattered_collate,
num_workers=1, # MUST be 1 or 0
drop_last=True,
pin_memory=True)
print('[{}] 1 training epoch = {} iters'.format(
args.epoch_iters,
datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
# create loader iterator
iterator_train = iter(loader_train)
# load nets into gpu
if args.num_gpus > 1:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=args.gpu_id)
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda(device=args.gpu_id[0])
# Set up optimizers
nets = (net_encoder, net_decoder, crit)
optimizers = create_optimizers(nets, args)
# Main loop
history = {
'train': {
'epoch': [],
'loss': [],
'train_acc': [],
'test_ious': [],
'test_mean_iou': []
}
}
for epoch in range(args.start_epoch, args.num_epoch + 1):
# test/validate
dataset_val = VOCValDataset(
args
) # create val dataset loader for every eval, in order to use drop_last=false
loader_val = torchdata.DataLoader(
dataset_val,
# data_parallel have been modified, MUST use val batch size
# and collate_fn MUST be user_scattered_collate
batch_size=args.val_batch_size,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=1, # MUST be 1 or 0
drop_last=False)
iterator_val = iter(loader_val)
if epoch % args.test_epoch_interval == 0 or epoch == args.num_epoch: # epoch != 1 and
(cls_ious, cls_mean_iou) = evaluate(segmentation_module,
iterator_val, args)
history['train']['test_ious'].append(cls_ious)
history['train']['test_mean_iou'].append(cls_mean_iou)
else:
history['train']['test_ious'].append(-1) # empty data
history['train']['test_mean_iou'].append(-1)
# train
train(segmentation_module, iterator_train, optimizers, history, epoch,
args)
# checkpointing
checkpoint(nets, history, args, epoch)
print('[{}] Training Done!'.format(
datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
def main():
"""Create the model and start the training."""
with open(args.config) as f:
config = yaml.load(f)
for k, v in config['common'].items():
setattr(args, k, v)
mkdirs(osp.join("logs/"+args.exp_name))
logger = create_logger('global_logger', "logs/" + args.exp_name + '/log.txt')
logger.info('{}'.format(args))
##############################
for key, val in vars(args).items():
logger.info("{:16} {}".format(key, val))
logger.info("random_scale {}".format(args.random_scale))
logger.info("is_training {}".format(args.is_training))
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
h, w = map(int, args.input_size_target.split(','))
input_size_target = (h, w)
print(type(input_size_target[1]))
cudnn.enabled = True
args.snapshot_dir = args.snapshot_dir + args.exp_name
tb_logger = SummaryWriter("logs/"+args.exp_name)
##############################
#validation data
local_array = np.load("local.npy")
local_array = local_array[:,:,:19]
local_array = local_array / local_array.sum(2).reshape(512, 1024, 1)
local_array = local_array.transpose(2,0,1)
local_array = torch.from_numpy(local_array)
local_array = local_array.view(1, 19, 512, 1024)
h, w = map(int, args.input_size_test.split(','))
input_size_test = (h,w)
h, w = map(int, args.com_size.split(','))
com_size = (h, w)
h, w = map(int, args.input_size_crop.split(','))
input_size_crop = h,w
h,w = map(int, args.input_size_target_crop.split(','))
input_size_target_crop = h,w
mean=[0.485, 0.456, 0.406]
std=[0.229, 0.224, 0.225]
normalize_module = transforms_seg.Normalize(mean=mean,
std=std)
test_normalize = transforms.Normalize(mean=mean,
std=std)
test_transform = transforms.Compose([
transforms.Resize((input_size_test[1], input_size_test[0])),
transforms.ToTensor(),
test_normalize])
valloader = data.DataLoader(cityscapesDataSet(
args.data_dir_target,
args.data_list_target_val,
crop_size=input_size_test,
set='train',
transform=test_transform),num_workers=args.num_workers,
batch_size=1, shuffle=False, pin_memory=True)
with open('./dataset/cityscapes_list/info.json', 'r') as fp:
info = json.load(fp)
mapping = np.array(info['label2train'], dtype=np.int)
label_path_list_val = args.label_path_list_val
gt_imgs_val = open(label_path_list_val, 'r').read().splitlines()
gt_imgs_val = [osp.join(args.data_dir_target_val, x) for x in gt_imgs_val]
name_classes = np.array(info['label'], dtype=np.str)
interp_val = nn.Upsample(size=(com_size[1], com_size[0]),mode='bilinear', align_corners=True)
####
#build model
####
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_classes,
weights=args.weights_decoder,
use_aux=True)
weighted_softmax = pd.read_csv("weighted_loss.txt", header=None)
weighted_softmax = weighted_softmax.values
weighted_softmax = torch.from_numpy(weighted_softmax)
weighted_softmax = weighted_softmax / torch.sum(weighted_softmax)
weighted_softmax = weighted_softmax.cuda().float()
model = SegmentationModule(
net_encoder, net_decoder, args.use_aux)
if args.num_gpus > 1:
model = torch.nn.DataParallel(model)
patch_replication_callback(model)
model.cuda()
nets = (net_encoder, net_decoder, None, None)
optimizers = create_optimizer(nets, args)
cudnn.enabled=True
cudnn.benchmark=True
model.train()
mean_mapping = [0.485, 0.456, 0.406]
mean_mapping = [item * 255 for item in mean_mapping]
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
source_transform = transforms_seg.Compose([
transforms_seg.Resize([input_size[1], input_size[0]]),
#segtransforms.RandScale((0.75, args.scale_max)),
#segtransforms.RandRotate((args.rotate_min, args.rotate_max), padding=mean_mapping, ignore_label=args.ignore_label),
#segtransforms.RandomGaussianBlur(),
#segtransforms.RandomHorizontalFlip(),
#segtransforms.Crop([input_size_crop[1], input_size_crop[0]], crop_type='rand', padding=mean_mapping, ignore_label=args.ignore_label),
transforms_seg.ToTensor(),
normalize_module])
target_transform = transforms_seg.Compose([
transforms_seg.Resize([input_size_target[1], input_size_target[0]]),
#segtransforms.RandScale((0.75, args.scale_max)),
#segtransforms.RandRotate((args.rotate_min, args.rotate_max), padding=mean_mapping, ignore_label=args.ignore_label),
#segtransforms.RandomGaussianBlur(),
#segtransforms.RandomHorizontalFlip(),
#segtransforms.Crop([input_size_target_crop[1], input_size_target_crop[0]],crop_type='rand', padding=mean_mapping, ignore_label=args.ignore_label),
transforms_seg.ToTensor(),
normalize_module])
trainloader = data.DataLoader(
GTA5DataSet(args.data_dir, args.data_list, max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size, transform = source_transform),
batch_size=args.batch_size, shuffle=True, num_workers=5, pin_memory=True)
trainloader_iter = enumerate(trainloader)
targetloader = data.DataLoader(fake_cityscapesDataSet(args.data_dir_target, args.data_list_target,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
set=args.set,
transform=target_transform),
batch_size=args.batch_size, shuffle=True, num_workers=5,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
# implement model.optim_parameters(args) to handle different models' lr setting
criterion_seg = torch.nn.CrossEntropyLoss(ignore_index=255,reduce=False)
criterion_pseudo = torch.nn.BCEWithLogitsLoss(reduce=False).cuda()
bce_loss = torch.nn.BCEWithLogitsLoss().cuda()
criterion_reconst = torch.nn.L1Loss().cuda()
criterion_soft_pseudo = torch.nn.MSELoss(reduce=False).cuda()
criterion_box = torch.nn.CrossEntropyLoss(ignore_index=255, reduce=False)
interp = nn.Upsample(size=(input_size[1], input_size[0]),align_corners=True, mode='bilinear')
interp_target = nn.Upsample(size=(input_size_target[1], input_size_target[0]), align_corners=True, mode='bilinear')
# labels for adversarial training
source_label = 0
target_label = 1
optimizer_encoder, optimizer_decoder, optimizer_disc, optimizer_reconst = optimizers
batch_time = AverageMeter(10)
loss_seg_value1 = AverageMeter(10)
best_mIoUs = 0
best_test_mIoUs = 0
loss_seg_value2 = AverageMeter(10)
loss_reconst_source_value = AverageMeter(10)
loss_reconst_target_value = AverageMeter(10)
loss_balance_value = AverageMeter(10)
loss_eq_att_value = AverageMeter(10)
loss_pseudo_value = AverageMeter(10)
bounding_num = AverageMeter(10)
pseudo_num = AverageMeter(10)
loss_bbx_att_value = AverageMeter(10)
for i_iter in range(args.num_steps):
# train G
# don't accumulate grads in D
end = time.time()
_, batch = trainloader_iter.__next__()
images, labels, _ = batch
images = Variable(images).cuda(async=True)
labels = Variable(labels).cuda(async=True)
results = model(images, labels)
loss_seg2 = results[-2]
loss_seg1 = results[-1]
loss_seg2 = torch.mean(loss_seg2)
loss_seg1 = torch.mean(loss_seg1)
loss = args.lambda_trade_off*(loss_seg2+args.lambda_seg * loss_seg1)
# proper normalization
#logger.info(loss_seg1.data.cpu().numpy())
loss_seg_value2.update(loss_seg2.data.cpu().numpy())
optimizer_encoder.zero_grad()
optimizer_decoder.zero_grad()
loss.backward()
optimizer_encoder.step()
optimizer_decoder.step()
_, batch = targetloader_iter.__next__()
images, fake_labels, _ = batch
images = Variable(images).cuda(async=True)
fake_labels = Variable(fake_labels, requires_grad=False).cuda()
results = model(images, None)
target_seg = results[0]
conf_tea, pseudo_label = torch.max(nn.functional.softmax(target_seg), dim=1)
pseudo_label = pseudo_label.detach()
# pseudo label hard
loss_pseudo = criterion_seg(target_seg, pseudo_label)
fake_mask = (fake_labels!=255).float().detach()
conf_mask = torch.gt(conf_tea, args.conf_threshold).float().detach()
loss_pseudo = loss_pseudo * conf_mask.detach() * fake_mask.detach()
loss_pseudo = loss_pseudo.view(-1)
loss_pseudo = loss_pseudo[loss_pseudo!=0]
#loss_pseudo = torch.sum(loss_pseudo * conf_mask.detach() * fake_mask.detach())
predict_class_mean = torch.mean(nn.functional.softmax(target_seg), dim=0).mean(1).mean(1)
equalise_cls_loss = robust_binary_crossentropy(predict_class_mean, weighted_softmax)
#equalise_cls_loss = torch.mean(equalise_cls_loss)* args.num_classes * torch.sum(conf_mask * fake_mask) / float(input_size_crop[0] * input_size_crop[1] * args.batch_size)
# new equalise_cls_loss
equalise_cls_loss = torch.mean(equalise_cls_loss)
#loss=args.lambda_balance * equalise_cls_loss
#bbx attention
loss_bbx_att = []
loss_eq_att = []
for box_idx, box_size in enumerate(args.box_size):
pooling = torch.nn.AvgPool2d(box_size)
pooling_result_i = pooling(target_seg)
local_i = pooling(local_array).float().cuda()
pooling_conf_mask, pooling_pseudo = torch.max(nn.functional.softmax(pooling_result_i), dim=1)
pooling_conf_mask = torch.gt(pooling_conf_mask, args.conf_threshold).float().detach()
fake_mask_i = pooling(fake_labels.unsqueeze(1).float())
fake_mask_i = fake_mask_i.squeeze(1)
fake_mask_i = (fake_mask_i!=255).float().detach()
loss_bbx_att_i = criterion_seg(pooling_result_i, pooling_pseudo)
loss_bbx_att_i = loss_bbx_att_i * pooling_conf_mask * fake_mask_i
loss_bbx_att_i = loss_bbx_att_i.view(-1)
loss_bbx_att_i = loss_bbx_att_i[loss_bbx_att_i!=0]
loss_bbx_att.append(loss_bbx_att_i)
pooling_result_i = pooling_result_i.mean(0).unsqueeze(0)
equalise_cls_loss_i = robust_binary_crossentropy(nn.functional.softmax(pooling_result_i), local_i)
equalise_cls_loss_i = equalise_cls_loss_i.mean(1)
equalise_cls_loss_i = equalise_cls_loss_i * pooling_conf_mask * fake_mask_i
equalise_cls_loss_i = equalise_cls_loss_i.view(-1)
equalise_cls_loss_i = equalise_cls_loss_i[equalise_cls_loss_i!=0]
loss_eq_att.append(equalise_cls_loss_i)
if len(args.box_size) > 0:
if args.merge_1x1:
loss_bbx_att.append(loss_pseudo)
loss_bbx_att = torch.cat(loss_bbx_att, dim=0)
bounding_num.update(loss_bbx_att.size(0) / float(560*480*args.batch_size))
loss_bbx_att = torch.mean(loss_bbx_att)
loss_eq_att = torch.cat(loss_eq_att, dim=0)
loss_eq_att = torch.mean(loss_eq_att)
loss_eq_att_value.update(loss_eq_att.item())
else:
loss_bbx_att = torch.mean(loss_pseudo)
loss_eq_att = 0
pseudo_num.update(loss_pseudo.size(0) / float(560*480*args.batch_size))
loss_pseudo = torch.mean(loss_pseudo)
if not args.merge_1x1:
loss += args.lambda_pseudo * loss_pseudo
loss = args.lambda_balance * equalise_cls_loss
if not isinstance(loss_bbx_att, list):
loss += args.lambda_pseudo * loss_bbx_att
loss += args.lambda_eq * loss_eq_att
loss_pseudo_value.update(loss_pseudo.item())
loss_balance_value.update(equalise_cls_loss.item())
optimizer_encoder.zero_grad()
optimizer_decoder.zero_grad()
loss.backward()
optimizer_encoder.step()
optimizer_decoder.step()
#optimizer_disc.step()
#loss_target_disc_value.update(loss_target_disc.data.cpu().numpy())
batch_time.update(time.time() - end)
remain_iter = args.num_steps - i_iter
remain_time = remain_iter * batch_time.avg
t_m, t_s = divmod(remain_time, 60)
t_h, t_m = divmod(t_m, 60)
remain_time = '{:02d}:{:02d}:{:02d}'.format(int(t_h), int(t_m), int(t_s))
if i_iter == args.decrease_lr:
adjust_learning_rate(optimizer_encoder, i_iter, args.lr_encoder, args)
adjust_learning_rate(optimizer_decoder, i_iter, args.lr_decoder, args)
if i_iter % args.print_freq == 0:
lr_encoder = optimizer_encoder.param_groups[0]['lr']
lr_decoder = optimizer_decoder.param_groups[0]['lr']
logger.info('exp = {}'.format(args.snapshot_dir))
logger.info('Iter = [{0}/{1}]\t'
'Time = {batch_time.avg:.3f}\t'
'loss_seg1 = {loss_seg1.avg:4f}\t'
'loss_seg2 = {loss_seg2.avg:.4f}\t'
'loss_reconst_source = {loss_reconst_source.avg:.4f}\t'
'loss_bbx_att = {loss_bbx_att.avg:.4f}\t'
'loss_reconst_target = {loss_reconst_target.avg:.4f}\t'
'loss_pseudo = {loss_pseudo.avg:.4f}\t'
'loss_eq_att = {loss_eq_att.avg:.4f}\t'
'loss_balance = {loss_balance.avg:.4f}\t'
'bounding_num = {bounding_num.avg:.4f}\t'
'pseudo_num = {pseudo_num.avg:4f}\t'
'lr_encoder = {lr_encoder:.8f} lr_decoder = {lr_decoder:.8f}'.format(
i_iter, args.num_steps, batch_time=batch_time,
loss_seg1=loss_seg_value1, loss_seg2=loss_seg_value2,
loss_pseudo=loss_pseudo_value,
loss_bbx_att = loss_bbx_att_value,
bounding_num = bounding_num,
loss_eq_att = loss_eq_att_value,
pseudo_num = pseudo_num,
loss_reconst_source=loss_reconst_source_value,
loss_balance=loss_balance_value,
loss_reconst_target=loss_reconst_target_value,
lr_encoder=lr_encoder,
lr_decoder=lr_decoder))
logger.info("remain_time: {}".format(remain_time))
if not tb_logger is None:
tb_logger.add_scalar('loss_seg_value1', loss_seg_value1.avg, i_iter)
tb_logger.add_scalar('loss_seg_value2', loss_seg_value2.avg, i_iter)
tb_logger.add_scalar('bounding_num', bounding_num.avg, i_iter)
tb_logger.add_scalar('pseudo_num', pseudo_num.avg, i_iter)
tb_logger.add_scalar('loss_pseudo', loss_pseudo_value.avg, i_iter)
tb_logger.add_scalar('lr', lr_encoder, i_iter)
tb_logger.add_scalar('loss_balance', loss_balance_value.avg, i_iter)
#####
#save image result
if i_iter % args.save_pred_every == 0 and i_iter != 0:
logger.info('taking snapshot ...')
model.eval()
val_time = time.time()
hist = np.zeros((19,19))
# f = open(args.result_dir, 'a')
# for index, batch in tqdm(enumerate(testloader)):
# with torch.no_grad():
# image, name = batch
# results = model(Variable(image).cuda(), None)
# output2 = results[0]
# pred = interp_val(output2)
# del output2
# pred = pred.cpu().data[0].numpy()
# pred = pred.transpose(1, 2, 0)
# pred = np.asarray(np.argmax(pred, axis=2), dtype=np.uint8)
# label = np.array(Image.open(gt_imgs_val[index]))
# #label = np.array(label.resize(com_size, Image.
# label = label_mapping(label, mapping)
# #logger.info(label.shape)
# hist += fast_hist(label.flatten(), pred.flatten(), 19)
# mIoUs = per_class_iu(hist)
# for ind_class in range(args.num_classes):
# logger.info('===>' + name_classes[ind_class] + ':\t' + str(round(mIoUs[ind_class] * 100, 2)))
# tb_logger.add_scalar(name_classes[ind_class] + '_mIoU', mIoUs[ind_class], i_iter)
# logger.info(mIoUs)
# tb_logger.add_scalar('val mIoU', mIoUs, i_iter)
# tb_logger.add_scalar('val mIoU', mIoUs, i_iter)
# f.write('i_iter:{:d},\tmiou:{:0.3f} \n'.format(i_iter, mIoUs))
# f.close()
# if mIoUs > best_mIoUs:
is_best = True
# best_mIoUs = mIoUs
#test validation
model.eval()
val_time = time.time()
hist = np.zeros((19,19))
# f = open(args.result_dir, 'a')
for index, batch in tqdm(enumerate(valloader)):
with torch.no_grad():
image, name = batch
results = model(Variable(image).cuda(), None)
output2 = results[0]
pred = interp_val(output2)
del output2
pred = pred.cpu().data[0].numpy()
pred = pred.transpose(1, 2, 0)
pred = np.asarray(np.argmax(pred, axis=2), dtype=np.uint8)
label = np.array(Image.open(gt_imgs_val[index]))
#label = np.array(label.resize(com_size, Image.
label = label_mapping(label, mapping)
#logger.info(label.shape)
hist += fast_hist(label.flatten(), pred.flatten(), 19)
mIoUs = per_class_iu(hist)
for ind_class in range(args.num_classes):
logger.info('===>' + name_classes[ind_class] + ':\t' + str(round(mIoUs[ind_class] * 100, 2)))
tb_logger.add_scalar(name_classes[ind_class] + '_mIoU', mIoUs[ind_class], i_iter)
mIoUs = round(np.nanmean(mIoUs) *100, 2)
is_best_test = False
logger.info(mIoUs)
tb_logger.add_scalar('test mIoU', mIoUs, i_iter)
if mIoUs > best_test_mIoUs:
best_test_mIoUs = mIoUs
is_best_test = True
# logger.info("best mIoU {}".format(best_mIoUs))
logger.info("best test mIoU {}".format(best_test_mIoUs))
net_encoder, net_decoder, net_disc, net_reconst = nets
save_checkpoint(net_encoder, 'encoder', i_iter, args, is_best_test)
save_checkpoint(net_decoder, 'decoder', i_iter, args, is_best_test)
is_best_test = False
model.train()
def main(args):
# 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=150,
weights=args.weights_decoder)
crit = nn.NLLLoss(ignore_index=-1)
if args.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit, args.deep_sup_scale)
else:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit)
########
for param in segmentation_module.encoder.parameters():
param.requires_grad = False
#for name, param in segmentation_module.decoder.named_parameters():
# print(name)
# if(name == "conv_last.weight" or name =="conv_last.bias" or name =="conv_last_deepsup.weight" or name =="conv_last_deepsup.bias"):
# param.requires_grad = True
#else:
# param.requires_grad = False
#print(param.requires_grad)
segmentation_module.decoder.conv_last = nn.Conv2d(args.fc_dim // 4, 12, 1,
1, 0)
#segmentation_module.decoder.conv_last.
segmentation_module.decoder.conv_last_deepsup = nn.Conv2d(
args.fc_dim // 4, 12, 1, 1, 0)
########
# Dataset and Loader
dataset_train = TrainDataset(args.list_train,
args,
batch_per_gpu=args.batch_size_per_gpu)
loader_train = torchdata.DataLoader(
dataset_train,
batch_size=len(args.gpus), # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=int(args.workers),
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(args.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
#######
#torch.backends.cudnn.benchmark = True
#CUDA_LAUNCH_BLOCKING=1
#######
# load nets into gpu
if len(args.gpus) > 1:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=args.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, args)
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
for epoch in range(args.start_epoch, args.num_epoch + 1):
train(segmentation_module, iterator_train, optimizers, history, epoch,
args)
# checkpointing
checkpoint(nets, history, args, epoch)
print('Training Done!')
def main(cfg, args, gpus):
USE_MULTI_GPU = True if len(gpus) > 1 else False
NUM_WORKERS = cfg.TEST.workers
if cfg.TEST.use_gpu:
BATCH_SIZE = cfg.TEST.batch_size_per_gpu * len(gpus)
else:
BATCH_SIZE = cfg.TEST.batch_size
RESULT_SAVE_DIR = cfg.TEST.result + '/' + cfg.MODEL.name + '/'
if not os.path.exists(RESULT_SAVE_DIR):
os.makedirs(RESULT_SAVE_DIR)
if not os.path.exists(RESULT_SAVE_DIR):
print("Error: Could not create result directory ", RESULT_SAVE_DIR)
exit()
test_set_dict = {
"test_airsim_city_0": [
1007, 1011, 1012, 1020, 1022, 1030, 1032, 1034, 1036, 2007, 2011,
2012, 2018, 2022, 2032, 2034, 2036, 3007, 3011, 3012, 3018, 3020,
3022, 3032, 3034, 3036, 4007, 4011, 4012, 4018, 4020, 4030, 4032,
4034, 4036
],
"test_jackal_0":
[5, 6, 7, 11, 15, 18, 19, 23, 24, 29, 30, 33, 37, 40, 43],
"test_jackal_t": [37],
"test_ahg_husky_t": [1]
}
session_list_test = [7, 9, 10]
if cfg.DATASET.test_set is not None:
session_list_test = test_set_dict[cfg.DATASET.test_set]
if cfg.TEST.use_gpu and torch.cuda.is_available():
device = torch.device("cuda:" + str(gpus[0]))
used_gpu_count = 1
total_mem = (
float(torch.cuda.get_device_properties(device).total_memory) /
1000000.0)
gpu_name = torch.cuda.get_device_name(device)
print("Using ", gpu_name, " with ", total_mem, " MB of memory.")
else:
device = torch.device("cpu")
used_gpu_count = 0
print("Output Model: ", cfg.MODEL.name)
print("Test data: ", session_list_test)
print("Network base model is ",
cfg.MODEL.arch_encoder + '+' + cfg.MODEL.arch_decoder)
print("Batch size: ", BATCH_SIZE)
print("Workers num: ", NUM_WORKERS)
print("device: ", device)
phases = ['test']
#data_set_portion_to_sample = {'train': 0.8, 'val': 0.2}
data_set_portion_to_sample = {'train': 1.0, 'val': 1.0}
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
input_img_width = int(cfg.DATASET.img_width)
input_img_height = int(cfg.DATASET.img_height)
target_img_width = int(cfg.DATASET.img_width)
target_img_height = int(cfg.DATASET.img_height)
# Transform loaded images. If not using color images, it will copy the single
# channel 3 times to keep the size of an RGB image.
if cfg.DATASET.use_color_images:
if cfg.DATASET.normalize_input:
data_transform_input = transforms.Compose([
transforms.Resize((input_img_height, input_img_width)),
transforms.ToTensor(), normalize
])
else:
data_transform_input = transforms.Compose([
transforms.Resize((input_img_height, input_img_width)),
transforms.ToTensor(),
])
else:
if cfg.DATASET.normalize_input:
data_transform_input = transforms.Compose([
transforms.Resize((input_img_height, input_img_width)),
transforms.ToTensor(),
transforms.Lambda(lambda x: torch.cat([x, x, x], 0)), normalize
])
else:
data_transform_input = transforms.Compose([
transforms.Resize((input_img_height, input_img_width)),
transforms.ToTensor(),
transforms.Lambda(lambda x: torch.cat([x, x, x], 0))
])
data_transform_target = transforms.Compose([
transforms.Resize((target_img_height, target_img_width)),
transforms.ToTensor()
])
# load_mask = cfg.TRAIN.use_masked_loss or cfg.MODEL.predict_conf_mask
load_mask = False
print("raw image root: " + cfg.DATASET.raw_img_root)
print("raw image folder: " + cfg.DATASET.raw_img_folder)
test_dataset = ImageQualityDataset(
cfg.DATASET.root,
cfg.DATASET.raw_img_root,
session_list_test,
loaded_image_color=cfg.DATASET.is_dataset_color,
output_image_color=cfg.DATASET.use_color_images,
session_prefix_length=cfg.DATASET.session_prefix_len,
raw_img_folder=cfg.DATASET.raw_img_folder,
no_meta_data_available=True,
transform_input=data_transform_input,
transform_target=data_transform_target,
load_masks=load_mask,
regression_mode=cfg.MODEL.is_regression_mode,
binarize_target=cfg.DATASET.binarize_target)
datasets = {phases[0]: test_dataset}
data_loaders = {
x: torch.utils.data.DataLoader(datasets[x],
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS)
for x in phases
}
# Build the network from selected modules
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,
regression_mode=cfg.MODEL.is_regression_mode,
inference_mode=True)
# The desired size of the output image. The model interpolates the output
# to this size
desired_size = (cfg.DATASET.img_height, cfg.DATASET.img_width)
# The desired size for the output of the network when saving it to file
# as an image
raw_output_img_size = (cfg.TEST.output_img_width,
cfg.TEST.output_img_height)
if cfg.MODEL.is_regression_mode:
if cfg.TRAIN.use_masked_loss:
criterion = MaskedMSELoss()
print("Regression Mode with Masked Loss")
else:
criterion = nn.MSELoss(reduction='mean')
print("Regression Mode")
else:
criterion = nn.NLLLoss(ignore_index=-1)
print("Segmentation Mode")
if cfg.MODEL.arch_decoder.endswith('deepsup'):
net = SegmentationModule(net_encoder,
net_decoder,
criterion,
cfg.TRAIN.deep_sup_scale,
segSize=desired_size)
else:
net = SegmentationModule(net_encoder,
net_decoder,
criterion,
segSize=desired_size)
if cfg.TEST.use_gpu and USE_MULTI_GPU:
if torch.cuda.device_count() >= len(gpus):
available_gpu_count = torch.cuda.device_count()
print("Using ", len(gpus), " GPUs out of available ",
available_gpu_count)
print("Used GPUs: ", gpus)
net = nn.DataParallel(net, device_ids=gpus)
# For synchronized batch normalization:
patch_replication_callback(net)
else:
print("Requested GPUs not available: ", gpus)
exit()
net = net.to(device)
print("Starting Inference...")
start_time = time.time()
# Runs inference on all data
for cur_phase in phases:
# Set model to evaluate mode
net.eval()
# Iterate over data
for i, data in enumerate(tqdm(data_loaders[cur_phase]), 0):
# get the inputs
input = data['img']
img_names = data['img_name']
session_nums = data['session']
feed_dict = dict()
feed_dict['input'] = input.to(device)
# Do not track history since we are in eval mode
with torch.set_grad_enabled(False):
# forward pass
output = net(feed_dict)
# output = torch.sigmoid(20 * (output - 0.5))
input_np = input.to(torch.device("cpu")).numpy()
output_np = output.to(torch.device("cpu")).numpy()
save_result_images(input_np,
target_imgs=None,
output_imgs=output_np,
img_names=img_names,
session_nums=session_nums,
save_dir=RESULT_SAVE_DIR,
raw_output_size=raw_output_img_size,
gt_available=False,
save_raw_output=cfg.TEST.save_raw_output,
initial_directory_prep=(i == 0))
if cur_phase == 'test':
if i % 100 == 99: # print every 100 mini-batches
print('[%5d]' % (i + 1))
if used_gpu_count:
print("Total GPU usage (MB): ",
calculate_gpu_usage(gpus), " / ",
used_gpu_count * total_mem)
print('All data was processed.')
time_elapsed = time.time() - start_time
print('Completed in {:.0f}h {:.0f}m {:.0f}s'.format(
time_elapsed // 3600, (time_elapsed % 3600) // 60, time_elapsed % 60))
def main(args):
# Network Builders
builder = ModelBuilder()
crit = nn.NLLLoss(ignore_index=-1)
crit = crit.cuda()
net_encoder = builder.build_encoder(
weights="baseline-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth")
gcu = GraphConv(
batch=args.batch_size_per_gpu
) #, V=2), GCU(X=enc_out, V=4), GCU(X=enc_out, V=8),GCU(X=enc_out, V=32)]
# gcu.load_state_dict(torch.load("ckpt/baseline-resnet50dilated-ngpus1-batchSize1-imgMaxSize1000-paddingConst8-segmDownsampleRate8-epoch20/decoder_epoch_20.pth"))
segmentation_module = SegmentationModule(net_encoder, gcu, crit, tr=True)
# Dataset and Loader
dataset_train = TrainDataset(args.list_train,
args,
batch_per_gpu=args.batch_size_per_gpu)
loader_train = torchdata.DataLoader(
dataset_train,
batch_size=len(args.gpus), # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=int(args.workers),
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(args.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
# load nets into gpu
if len(args.gpus) > 4:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=args.gpus)
# For sync bn
patch_replication_callback(segmentation_module)
# segmentation_module.cuda()
# Set up optimizers
# print(gcu[0].parameters())
nets = (net_encoder, gcu, crit)
optimizers, par = create_optimizers(nets, args)
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
vis = visdom.Visdom()
win = vis.line(np.array([5.7]),
opts=dict(xlabel='epochs',
ylabel='Loss',
title='Training Loss V=16',
legend=['Loss']))
for epoch in range(args.start_epoch, args.num_epoch + 1):
lss = train(segmentation_module, iterator_train, optimizers, history,
epoch, par, vis, win, args)
# checkpointing
checkpoint(nets, history, args, epoch)
print('Training Done!')
def main(cfg, args, gpus):
USE_MULTI_GPU = True if len(gpus) > 1 else False
NUM_WORKERS = cfg.TRAIN.workers
if cfg.TRAIN.use_gpu:
BATCH_SIZE = cfg.TRAIN.batch_size_per_gpu * len(gpus)
else:
BATCH_SIZE = cfg.TRAIN.batch_size
train_set_dict = {
"train_airsim_city_0": [
1005, 1006, 1009, 1010, 1013, 1014, 1015, 1016, 1021, 1023, 1025,
1027, 1028, 1029, 1031, 1035, 2005, 2006, 2009, 2010, 2013, 2014,
2015, 2016, 2021, 2023, 2025, 2027, 2028, 2029, 2031, 2033, 2035,
3005, 3006, 3009, 3010, 3013, 3014, 3015, 3016, 3021, 3023, 3025,
3027, 3028, 3029, 3031, 3033, 3035, 4005, 4006, 4009, 4010, 4013,
4014, 4015, 4021, 4023, 4025, 4027, 4028, 4029, 4031, 4033, 4035
],
"train_jackal_0":
[1, 2, 3, 4, 8, 10, 13, 16, 17, 20, 22, 25, 27, 28, 31, 36, 42],
"train_jackal_t": [37],
"train_stereo_2020_12_21_run1_t": [1],
"train_stereo_2021_01_13_run2_t": [1]
}
valid_set_dict = {
"valid_airsim_city_0": [1008, 1016, 1017, 1024],
"valid_jackal_0": [14],
"valid_jackal_t": [37],
"valid_stereo_2020_12_21_run1_t": [1],
"valid_stereo_2021_01_13_run2_t": [1]
}
session_list_train = [0, 1, 2, 3, 4, 5]
session_list_val = [6, 8]
session_list_test = [7, 9, 10]
if (cfg.DATASET.train_set is not None) and (cfg.DATASET.validation_set
is not None):
session_list_train = train_set_dict[cfg.DATASET.train_set]
session_list_val = valid_set_dict[cfg.DATASET.validation_set]
if cfg.TRAIN.use_gpu and torch.cuda.is_available():
device = torch.device("cuda:" + str(gpus[0]))
used_gpu_count = 1
total_mem = (
float(torch.cuda.get_device_properties(device).total_memory) /
1000000.0)
gpu_name = torch.cuda.get_device_name(device)
print("Using ", gpu_name, " with ", total_mem, " MB of memory.")
else:
device = torch.device("cpu")
used_gpu_count = 0
print("Output Model: ", cfg.MODEL.name)
print("Train data: ", session_list_train)
print("Validation data: ", session_list_val)
print("Network base model is ",
cfg.MODEL.arch_encoder + '+' + cfg.MODEL.arch_decoder)
print("Batch size: ", BATCH_SIZE)
print("Workers num: ", NUM_WORKERS)
print("device: ", device)
phases = ['train', 'val']
#data_set_portion_to_sample = {'train': 0.8, 'val': 0.2}
data_set_portion_to_sample = {'train': 1.0, 'val': 1.0}
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
input_img_width = int(cfg.DATASET.img_width)
input_img_height = int(cfg.DATASET.img_height)
target_img_width = int(cfg.DATASET.img_width /
cfg.DATASET.target_downsampling_rate)
target_img_height = int(cfg.DATASET.img_height /
cfg.DATASET.target_downsampling_rate)
# Transform loaded images. If not using color images, it will copy the single
# channel 3 times to keep the size of an RGB image.
if cfg.DATASET.use_color_images:
if cfg.DATASET.normalize_input:
data_transform_input = transforms.Compose([
transforms.Resize((input_img_height, input_img_width)),
transforms.ToTensor(), normalize
])
else:
data_transform_input = transforms.Compose([
transforms.Resize((input_img_height, input_img_width)),
transforms.ToTensor(),
])
else:
if cfg.DATASET.normalize_input:
data_transform_input = transforms.Compose([
transforms.Resize((input_img_height, input_img_width)),
transforms.ToTensor(),
transforms.Lambda(lambda x: torch.cat([x, x, x], 0)), normalize
])
else:
data_transform_input = transforms.Compose([
transforms.Resize((input_img_height, input_img_width)),
transforms.ToTensor(),
transforms.Lambda(lambda x: torch.cat([x, x, x], 0))
])
data_transform_target = transforms.Compose([
transforms.Resize((target_img_height, target_img_width)),
transforms.ToTensor()
])
load_mask = cfg.TRAIN.use_masked_loss or cfg.MODEL.predict_conf_mask
train_dataset = ImageQualityDataset(
cfg.DATASET.root,
cfg.DATASET.raw_img_root,
session_list_train,
loaded_image_color=cfg.DATASET.is_dataset_color,
output_image_color=cfg.DATASET.use_color_images,
session_prefix_length=cfg.DATASET.session_prefix_len,
raw_img_folder=cfg.DATASET.raw_img_folder,
no_meta_data_available=True,
load_only_with_labels=True,
transform_input=data_transform_input,
transform_target=data_transform_target,
load_masks=load_mask,
regression_mode=cfg.MODEL.is_regression_mode,
binarize_target=cfg.DATASET.binarize_target)
val_dataset = ImageQualityDataset(
cfg.DATASET.root,
cfg.DATASET.raw_img_root,
session_list_val,
loaded_image_color=cfg.DATASET.is_dataset_color,
output_image_color=cfg.DATASET.use_color_images,
session_prefix_length=cfg.DATASET.session_prefix_len,
raw_img_folder=cfg.DATASET.raw_img_folder,
no_meta_data_available=True,
load_only_with_labels=True,
transform_input=data_transform_input,
transform_target=data_transform_target,
load_masks=load_mask,
regression_mode=cfg.MODEL.is_regression_mode,
binarize_target=cfg.DATASET.binarize_target)
datasets = {phases[0]: train_dataset, phases[1]: val_dataset}
data_loaders = {
x: torch.utils.data.DataLoader(datasets[x],
batch_size=BATCH_SIZE,
num_workers=NUM_WORKERS)
for x in phases
}
# Build the network from selected modules
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,
regression_mode=cfg.MODEL.is_regression_mode,
inference_mode=False)
if cfg.MODEL.is_regression_mode:
if cfg.TRAIN.use_masked_loss:
criterion = MaskedMSELoss()
print("Regression Mode with Masked Loss")
else:
criterion = nn.MSELoss(reduction='mean')
print("Regression Mode")
else:
criterion = nn.NLLLoss(ignore_index=-1)
print("Segmentation Mode")
use_mask = cfg.TRAIN.use_masked_loss and cfg.MODEL.is_regression_mode
if cfg.MODEL.arch_decoder.endswith('deepsup'):
net = SegmentationModule(net_encoder,
net_decoder,
criterion,
cfg.TRAIN.deep_sup_scale,
use_mask=use_mask)
else:
net = SegmentationModule(net_encoder,
net_decoder,
criterion,
use_mask=use_mask)
if cfg.TRAIN.use_gpu and USE_MULTI_GPU:
if torch.cuda.device_count() >= len(gpus):
available_gpu_count = torch.cuda.device_count()
print("Using ", len(gpus), " GPUs out of available ",
available_gpu_count)
print("Used GPUs: ", gpus)
net = nn.DataParallel(net, device_ids=gpus)
# For synchronized batch normalization:
patch_replication_callback(net)
else:
print("Requested GPUs not available: ", gpus)
exit()
net = net.to(device)
# Set up optimizers
modules = (net_encoder, net_decoder, criterion)
optimizers = create_optimizers(modules, cfg)
best_loss = 1000000000.0
best_model_enc = copy.deepcopy(net_encoder.state_dict())
best_model_dec = copy.deepcopy(net_decoder.state_dict())
training_history = {x: {'loss': [], 'acc': []} for x in phases}
print("Starting Training...")
start_time = time.time()
# Runs training and validation
for epoch in range(cfg.TRAIN.num_epoch):
for cur_phase in phases:
if cur_phase == 'train':
net.train() # Set model to training mode
else:
net.eval() # Set model to evaluate mode
epoch_loss = 0.0
running_loss = 0.0
# Iterate over data
i = 0
for i, data in enumerate(data_loaders[cur_phase], 0):
# get the inputs
input = data['img']
feed_dict = dict()
feed_dict['input'] = input.to(device)
if not cfg.MODEL.is_regression_mode:
target = data['labels']
feed_dict['target'] = target.to(device)
else:
target = data['score_img']
feed_dict['target'] = target.to(device)
if cfg.TRAIN.use_masked_loss:
mask = data['mask_img']
feed_dict['mask'] = mask.to(device)
# zero the parameter gradients
net.zero_grad()
# track history if only in train
with torch.set_grad_enabled(cur_phase == 'train'):
# forward + backward + optimize
loss, acc = net(feed_dict)
# For multi gpu case, loss will be a vector the same length
# as the number of gpus. This is a side effect of loss function
# being inside the module.
loss = loss.mean()
# backward and optimize only if in training phase
if cur_phase == 'train':
loss.backward()
for optimizer in optimizers:
optimizer.step()
# print statistics
loss = loss.item()
running_loss += loss
epoch_loss += loss
if cur_phase == 'train':
if i % 100 == 99: # print every 100 mini-batches
print('[%d, %5d] Loss: %.6f' %
(epoch + 1, i + 1, running_loss / 100))
running_loss = 0.0
if used_gpu_count:
print("Total GPU usage (MB): ",
calculate_gpu_usage(gpus), " / ",
used_gpu_count * total_mem)
epoch_loss = epoch_loss / i
print('%s: Loss: %.6f' % (cur_phase, epoch_loss))
# Keep the best model so far
if cur_phase == 'val' and epoch_loss < best_loss:
best_loss = epoch_loss
best_model_enc = copy.deepcopy(net_encoder.state_dict())
best_model_dec = copy.deepcopy(net_decoder.state_dict())
training_history[cur_phase]['loss'].append(epoch_loss)
print('Epoch #%d finished. *******************' % (epoch + 1))
if (epoch + 1) % cfg.TRAIN.snapshot_interval == 0:
last_model_state = (net_encoder.state_dict(),
net_decoder.state_dict())
best_model_state = (best_model_enc, best_model_dec)
save_results(cfg.DIR,
cfg.MODEL.name,
last_model_state,
best_model_state,
training_history,
snapshot_idx=epoch)
print('Snapshot saved.')
print('Finished Training')
time_elapsed = time.time() - start_time
print('Completed in {:.0f}h {:.0f}m {:.0f}s'.format(
time_elapsed // 3600, (time_elapsed % 3600) // 60, time_elapsed % 60))
last_model_state = (net_encoder.state_dict(), net_decoder.state_dict())
best_model_state = (best_model_enc, best_model_dec)
save_results(cfg.DIR,
cfg.MODEL.name,
last_model_state,
best_model_state,
training_history,
snapshot_idx=-1)
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, gpus):
# Network Builders
net_encoder = models.ModelBuilder.build_encoder(
arch=cfg.MODEL.arch_encoder.lower(),
fc_dim=cfg.MODEL.fc_dim,
weights=cfg.MODEL.weights_encoder)
net_decoder = models.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)
crit = nn.NLLLoss(ignore_index=-1)
# crit = nn.CrossentropyLoss()
if cfg.MODEL.arch_decoder.endswith('deepsup'):
segmentation_module = models.SegmentationModule(
net_encoder, net_decoder, crit, None)
else:
segmentation_module = models.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), # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=cfg.TRAIN.workers,
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(cfg.TRAIN.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
# 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': []}}
for epoch in range(cfg.TRAIN.start_epoch, cfg.TRAIN.num_epoch):
train(segmentation_module, iterator_train, optimizers, history,
epoch + 1, cfg)
# checkpointing
checkpoint(nets, history, cfg, epoch + 1)
print('Training Done!')
def main(args):
# 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)
else:
unet = builder.build_unet(num_class=args.num_class,
arch=args.unet_arch,
weights=args.weights_unet)
print("Froze the following layers: ")
for name, p in unet.named_parameters():
if p.requires_grad == False:
print(name)
crit = nn.NLLLoss()
#crit = nn.BCEWithLogitsLoss(pos_weight=torch.tensor(50))
#crit = nn.CrossEntropyLoss().cuda()
#crit = nn.BCELoss()
if args.arch_decoder.endswith('deepsup') and args.unet == False:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit, args.deep_sup_scale)
else:
segmentation_module = SegmentationModule(net_encoder,
net_decoder,
crit,
is_unet=args.unet,
unet=unet)
train_augs = Compose([
RandomSized(224),
RandomHorizontallyFlip(),
RandomVerticallyFlip(),
RandomRotate(180),
AdjustContrast(cf=0.25),
AdjustBrightness(bf=0.25)
]) #, RandomErasing()])
#train_augs = None
# Dataset and Loader
dataset_train = TrainDataset(args.list_train,
args,
batch_per_gpu=args.batch_size_per_gpu,
augmentations=train_augs)
loader_train = data.DataLoader(
dataset_train,
batch_size=len(args.gpus), # we have modified data_parallel
shuffle=False, # we do not use this param
num_workers=int(args.workers),
drop_last=True,
pin_memory=False)
print('1 Epoch = {} iters'.format(args.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
# load nets into gpu
if len(args.gpus) > 1:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=args.gpus)
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder, crit) if args.unet == False else (unet,
crit)
optimizers = create_optimizers(nets, args)
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
for epoch in range(args.start_epoch, args.num_epoch + 1):
train(segmentation_module, iterator_train, optimizers, history, epoch,
args)
# checkpointing
checkpoint(nets, history, args, epoch)
print('Training Done!')
def main(args):
# Network Builders
builder = ModelBuilder()
unet = builder.build_unet(num_class=args.num_class,
arch=args.unet_arch,
weights=args.weights_unet)
print("Froze the following layers: ")
for name, p in unet.named_parameters():
if p.requires_grad == False:
print(name)
print()
crit = DualLoss(mode="train")
segmentation_module = SegmentationModule(crit, unet)
test_augs = Compose([PaddingCenterCrop(256)])
print("ready to load data")
dataset_val = LungData(
root=args.data_root,
split='test',
k_split=args.k_split,
augmentations=test_augs)
loader_val = data.DataLoader(
dataset_val,
batch_size=1,
shuffle=False,
collate_fn=user_scattered_collate,
num_workers=5,
drop_last=True)
print(len(loader_val))
# load nets into gpu
if len(args.gpus) > 1:
segmentation_module = UserScatteredDataParallel(
segmentation_module,
device_ids=args.gpus)
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder, crit) if args.unet == False else (unet, crit)
optimizers = create_optimizers(nets, args)
'''
# Start the webapp: user update a dcm file, output the predicted segmentation pic of it
inp = gradio.inputs.DcmUpload(preprocessing_fn=preprocess)
#inp = gradio.inputs.ImageUpload(preprocessing_fn=preprocess)
io = gradio.Interface(inputs=inp, outputs="image", model_type="lung_seg", model=segmentation_module, args=args)
io.launch(validate=False)
'''
iou, loss = eval(loader_val, segmentation_module, args, crit)
print('Evaluation Done!')
def main():
args = get_arguments()
# seeding for reproducbility
if torch.cuda.is_available():
torch.cuda.manual_seed(args.random_seed)
torch.manual_seed(args.random_seed)
np.random.seed(args.random_seed)
# instantiate dataset
dataset = AdobeImageMattingDataset
snapshot_dir = os.path.join(args.snapshot_dir, args.dataset.lower(),
args.exp)
if not os.path.exists(snapshot_dir):
os.makedirs(snapshot_dir)
args.result_dir = os.path.join(args.result_dir, args.exp)
if not os.path.exists(args.result_dir):
os.makedirs(args.result_dir)
args.restore_from = os.path.join(args.snapshot_dir, args.dataset.lower(),
args.exp, args.restore_from)
arguments = vars(args)
for item in arguments:
print(item, ':\t', arguments[item])
# instantiate network
hlnet = hlbackbone[args.backbone]
net = hlnet(pretrained=True,
freeze_bn=True,
output_stride=args.output_stride,
input_size=args.crop_size,
apply_aspp=args.apply_aspp,
conv_operator=args.conv_operator,
decoder=args.decoder,
decoder_kernel_size=args.decoder_kernel_size,
indexnet=args.indexnet,
index_mode=args.index_mode,
use_nonlinear=args.use_nonlinear,
use_context=args.use_context,
sync_bn=args.sync_bn)
if args.backbone == 'mobilenetv2':
net = nn.DataParallel(net)
if args.sync_bn:
patch_replication_callback(net)
net.cuda()
# filter parameters
pretrained_params = []
learning_params = []
for p in net.named_parameters():
if 'dconv' in p[0] or 'pred' in p[0] or 'index' in p[0]:
learning_params.append(p[1])
else:
pretrained_params.append(p[1])
# define optimizer
optimizer = torch.optim.Adam([
{
'params': learning_params
},
{
'params': pretrained_params,
'lr': args.learning_rate / args.mult
},
],
lr=args.learning_rate)
# restore parameters
start_epoch = 0
net.train_loss = {'running_loss': [], 'epoch_loss': []}
net.val_loss = {'running_loss': [], 'epoch_loss': []}
net.measure = {'sad': [], 'mse': [], 'grad': [], 'conn': []}
if args.restore_from is not None:
if os.path.isfile(args.restore_from):
checkpoint = torch.load(args.restore_from)
net.load_state_dict(checkpoint['state_dict'])
if 'epoch' in checkpoint:
start_epoch = checkpoint['epoch']
if 'optimizer' in checkpoint:
optimizer.load_state_dict(checkpoint['optimizer'])
if 'train_loss' in checkpoint:
net.train_loss = checkpoint['train_loss']
if 'val_loss' in checkpoint:
net.val_loss = checkpoint['val_loss']
if 'measure' in checkpoint:
net.measure = checkpoint['measure']
print("==> load checkpoint '{}' (epoch {})".format(
args.restore_from, start_epoch))
else:
with open(os.path.join(args.result_dir, args.exp + '.txt'),
'a') as f:
for item in arguments:
print(item, ':\t', arguments[item], file=f)
print("==> no checkpoint found at '{}'".format(args.restore_from))
# define transform
transform_train_val = [
RandomCrop(args.crop_size, args.scales),
RandomFlip()
]
transform_all = [
Normalize(args.image_scale, args.image_mean, args.image_std),
ToTensor()
]
composed_transform_train = transforms.Compose(transform_train_val +
transform_all)
composed_transform_val = transforms.Compose(transform_all)
# define dataset loader
trainset = dataset(data_file=args.data_list,
data_dir=args.data_dir,
train=True,
transform=composed_transform_train)
train_loader = DataLoader(trainset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True)
valset = dataset(data_file=args.data_val_list,
data_dir=args.data_dir,
train=False,
transform=composed_transform_val)
val_loader = DataLoader(valset,
batch_size=1,
shuffle=False,
num_workers=0,
pin_memory=True)
print('alchemy start...')
if args.evaluate_only:
validate(net, val_loader, start_epoch + 1, args)
return
resume_epoch = -1 if start_epoch == 0 else start_epoch
scheduler = MultiStepLR(optimizer,
milestones=[20, 26],
gamma=0.1,
last_epoch=resume_epoch)
for epoch in range(start_epoch, args.num_epochs):
scheduler.step()
# train
train(net, train_loader, optimizer, epoch + 1, scheduler, args)
# val
validate(net, val_loader, epoch + 1, args)
# save checkpoint
state = {
'state_dict': net.state_dict(),
'optimizer': optimizer.state_dict(),
'epoch': epoch + 1,
'train_loss': net.train_loss,
'val_loss': net.val_loss,
'measure': net.measure
}
save_checkpoint(state, snapshot_dir, filename='model_ckpt.pth.tar')
print(args.exp + ' epoch {} finished!'.format(epoch + 1))
if len(net.measure['grad']) > 1 and net.measure['grad'][-1] <= min(
net.measure['grad'][:-1]):
save_checkpoint(state, snapshot_dir, filename='model_best.pth.tar')
print('Experiments with ' + args.exp + ' done!')
def main(args):
# 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)
crit = nn.NLLLoss(ignore_index=-1)
if args.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit, args.deep_sup_scale)
else:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit)
# Dataset and Loader
dataset_train = TrainDataset(args.list_train,
args,
batch_per_gpu=args.batch_size_per_gpu)
loader_train = torchdata.DataLoader(
dataset_train,
batch_size=args.num_gpus, # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=int(args.workers),
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(args.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
# load nets into gpu
if args.num_gpus > 1:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=range(
args.num_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, args)
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
for epoch in range(args.start_epoch, args.num_epoch + 1):
train(segmentation_module, iterator_train, optimizers, history, epoch,
args)
# checkpointing
checkpoint(nets, history, args, epoch)
print('Training Done!')
def main(args):
# 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)
crit = nn.NLLLoss(ignore_index=255)
if args.arch_decoder.endswith('deepsup'):
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit, args.deep_sup_scale)
else:
segmentation_module = SegmentationModule(net_encoder, net_decoder,
crit)
# Dataset and Loader
# dataset_train = TrainDataset(
# args.list_train, args, batch_per_gpu=args.batch_size_per_gpu)
# dataset_train = voc.TrainDataset_VOC(dataset_root=config.img_root_folder, mode='train', transform=input_transform)
dataset_train = VOC_TrainDataset(opt=args,
batch_per_gpu=args.batch_size_per_gpu)
loader_train = torchdata.DataLoader(
dataset_train,
batch_size=args.num_gpus, # we have modified data_parallel
shuffle=False, # we do not use this param
collate_fn=user_scattered_collate,
num_workers=int(args.workers),
drop_last=True,
pin_memory=True)
print('1 Epoch = {} iters'.format(args.epoch_iters))
# create loader iterator
iterator_train = iter(loader_train)
# load nets into gpu
if args.num_gpus > 1:
segmentation_module = UserScatteredDataParallel(segmentation_module,
device_ids=range(
args.num_gpus))
# segmentation_module = UserScatteredDataParallel(
# segmentation_module,
# device_ids=[0,3,4,5])
# For sync bn
patch_replication_callback(segmentation_module)
segmentation_module.cuda()
# Set up optimizers
nets = (net_encoder, net_decoder, crit)
optimizers = create_optimizers(nets, args)
if args.resume:
file_path_history = '{}/history_{}'.format(args.ckpt, 'epoch_last.pth')
file_path_encoder = '{}/encoder_{}'.format(args.ckpt, 'epoch_last.pth')
file_path_decoder = '{}/decoder_{}'.format(args.ckpt, 'epoch_last.pth')
file_path_optimizers = '{}/optimizers_{}'.format(
args.ckpt, 'epoch_last.pth')
if os.path.isfile(file_path_history):
print("=> loading checkpoint '{}'".format(file_path_history))
checkpoint_history = torch.load(file_path_history)
checkpoint_encoder = torch.load(file_path_encoder)
checkpoint_decoder = torch.load(file_path_decoder)
checkpoint_optimizers = torch.load(file_path_optimizers)
args.start_epoch = int(checkpoint_history['train']['epoch'][0]) + 1
nets[0].load_state_dict(checkpoint_encoder)
nets[1].load_state_dict(checkpoint_decoder)
optimizers[0].load_state_dict(
checkpoint_optimizers['encoder_optimizer'])
optimizers[1].load_state_dict(
checkpoint_optimizers['decoder_optimizer'])
print("=> start train epoch {}".format(args.start_epoch))
else:
print('resume not find epoch-last checkpoint')
# Main loop
history = {'train': {'epoch': [], 'loss': [], 'acc': []}}
for epoch in range(args.start_epoch, args.num_epoch + 1):
train(segmentation_module, iterator_train, optimizers, history, epoch,
args)
# checkpointing
checkpoint(nets, optimizers, history, args, epoch)
print('Training Done!')
