changed some stuff, re-added default CREStereo scheduler

main
Cpt.Captain 2 years ago
parent 37c537ca31
commit 2731ef1ada
  1. 25
      cfgs/train.yaml
  2. 71
      dataset.py
  3. 172
      train_lightning.py

@ -1,12 +1,12 @@
seed: 0
mixed_precision: false
# base_lr: 4.0e-4
base_lr: 0.001
t_max: 161
mixed_precision: true
base_lr: 4.0e-4
# base_lr: 0.00001
t_max: 16100
nr_gpus: 3
batch_size: 2
n_total_epoch: 300
batch_size: 3
n_total_epoch: 100
minibatch_per_epoch: 500
loadmodel: ~
@ -17,11 +17,18 @@ model_save_freq_epoch: 1
max_disp: 256
image_width: 640
image_height: 480
# dataset: "blender"
# training_data_path: "./stereo_trainset/crestereo"
pattern_attention: false
dataset: "blender"
# training_data_path: "/media/Data1/connecting_the_dots_data/ctd_data/"
training_data_path: "/media/Data1/connecting_the_dots_data/blender_renders/data"
# training_data_path: "/media/Data1/connecting_the_dots_data/blender_renders/data"
training_data_path: "/media/Data1/connecting_the_dots_data/blender_renders_ctd_randomize_light/data"
# FIXME any of this??
pattern_attention: false
scene_attention: true
ignore_pattern_completely: false
log_level: "logging.INFO"

@ -7,6 +7,11 @@ from PIL import Image, ImageEnhance
from megengine.data.dataset import Dataset
def downsample(img):
downsampled = cv2.pyrDown(img)
diff = (downsampled.shape[0] - 480) // 2
return downsampled[diff:downsampled.shape[0]-diff, 0:downsampled.shape[1]]
class Augmentor:
def __init__(
@ -234,11 +239,12 @@ class CREStereoDataset(Dataset):
class CTDDataset(Dataset):
def __init__(self, root, pattern_path: str, data_type: str = 'syn', augment=True, resize_pattern=True, blur=False, split=0.9, test_set=False):
def __init__(self, root, pattern_path: str, data_type: str = 'syn', augment=True, resize_pattern=True, blur=False, split=0.9, test_set=False, use_lightning=True):
super().__init__()
self.rng = np.random.RandomState(0)
self.augment = augment
self.blur = blur
self.use_lightning = use_lightning
imgs = glob.glob(os.path.join(root, f"{data_type}/*/im0_*.npy"), recursive=True)
if test_set:
self.imgs = imgs[:int(split * len(imgs))]
@ -248,10 +254,7 @@ class CTDDataset(Dataset):
if resize_pattern and self.pattern.shape != (480, 640, 3):
# self.pattern = cv2.resize(self.pattern, (640, 480))
print(self.pattern.shape)
downsampled = cv2.pyrDown(self.pattern)
diff = (downsampled.shape[0] - 480) // 2
self.pattern = downsampled[diff:downsampled.shape[0]-diff, 0:downsampled.shape[1]]
self.pattern = downsample(self.pattern)
self.augmentor = Augmentor(
image_height=480,
@ -304,14 +307,18 @@ class CTDDataset(Dataset):
left_img, right_img, left_disp
)
right_img = right_img.transpose((2, 0, 1)).astype("uint8")
return {
"left": left_img,
"right": right_img,
"disparity": left_disp,
"mask": disp_mask,
}
if not self.use_lightning:
right_img = right_img.transpose((2, 0, 1)).astype("uint8")
return {
"left": left_img,
"right": right_img,
"disparity": left_disp,
"mask": disp_mask,
}
right_img = right_img.transpose((2, 0, 1)).astype("uint8")
left_img = left_img.transpose((2, 0, 1)).astype("uint8")
return left_img, right_img, left_disp, disp_mask
def __len__(self):
return len(self.imgs)
@ -321,17 +328,30 @@ class BlenderDataset(CTDDataset):
def __init__(self, root, pattern_path: str, data_type: str = 'syn', augment=True, resize_pattern=True, blur=False, split=0.9, test_set=False, use_lightning=False):
super().__init__(root, pattern_path)
self.use_lightning = use_lightning
imgs = [f for f in glob.glob(f"{root}/im_*.png", recursive=True) if not 'depth0001' in f]
if test_set:
additional_img_types = {
'depth',
'disp',
'grad',
}
pngs = glob.glob(f"{root}/im_*.png", recursive=True)
imgs = [
img for img in pngs
if all(
map(
lambda x: x not in img, additional_img_types
)
)
]
if not test_set:
self.imgs = imgs[:int(split * len(imgs))]
else:
self.imgs = imgs[int(split * len(imgs)):]
self.pattern = cv2.imread(pattern_path)#, cv2.IMREAD_GRAYSCALE)
if resize_pattern and self.pattern.shape != (480, 640, 3):
downsampled = cv2.pyrDown(self.pattern)
diff = (downsampled.shape[0] - 480) // 2
self.pattern = downsampled[diff:downsampled.shape[0]-diff, 0:downsampled.shape[1]]
self.pattern = downsample(self.pattern)
self.augmentor = Augmentor(
image_height=480,
@ -345,7 +365,7 @@ class BlenderDataset(CTDDataset):
def __getitem__(self, index):
# find path
left_path = self.imgs[index]
left_disp_path = left_path.split('.')[0] + '_depth0001.png'
left_disp_path = left_path.split('.')[0] + '_disp0001.png'
# read img, disp
left_img = cv2.imread(left_path)
@ -354,14 +374,14 @@ class BlenderDataset(CTDDataset):
left_img = (left_img * 255).astype('uint8')
if left_img.shape != (480, 640, 3):
downsampled = cv2.pyrDown(left_img)
diff = (downsampled.shape[0] - 480) // 2
left_img = downsampled[diff:downsampled.shape[0]-diff, 0:downsampled.shape[1]]
left_img = downsample(left_img)
if left_img.shape[-1] != 3:
left_img = cv2.merge([left_img, left_img, left_img]).reshape((480, 640, 3))
right_img = self.pattern
left_disp = self.get_disp(left_disp_path)
# left_disp = self.get_disp(left_disp_path)
disp = cv2.imread(left_disp_path, cv2.IMREAD_UNCHANGED)
left_disp = downsample(disp)
if False: # self.rng.binomial(1, 0.5):
left_img, right_img = np.fliplr(right_img), np.fliplr(left_img)
@ -384,7 +404,6 @@ class BlenderDataset(CTDDataset):
)
if not self.use_lightning:
# right_img = right_img.transpose((2, 0, 1)).astype("uint8")
return {
"left": left_img,
"right": right_img,
@ -400,15 +419,13 @@ class BlenderDataset(CTDDataset):
baseline = 0.075 # meters
fl = 560. # as per CTD
depth = cv2.imread(path, cv2.IMREAD_UNCHANGED)
downsampled = cv2.pyrDown(depth)
diff = (downsampled.shape[0] - 480) // 2
depth = downsampled[diff:downsampled.shape[0]-diff, 0:downsampled.shape[1]]
depth = downsample(depth)
# disp = np.load(path).transpose(1,2,0)
# disp = baseline * fl / depth
# return disp.astype(np.float32) / 32
# FIXME temporarily increase disparity until new data with better depth values is generated
# higher values seem to speedup convergence, but introduce much stronger artifacting
mystery_factor = 150
mystery_factor = 35
# mystery_factor = 1
disp = (baseline * fl * mystery_factor) / depth
return disp.astype(np.float32)

@ -94,6 +94,32 @@ def format_time(elapse):
return "{:02d}:{:02d}:{:02d}".format(hour, minute, seconds)
def outlier_fraction(estimate, target, mask=None, threshold=0):
def _process_inputs(estimate, target, mask):
if estimate.shape != target.shape:
raise Exception(f'estimate and target have to be same shape (expected {estimate.shape} == {target.shape})')
if mask is None:
mask = np.ones(estimate.shape, dtype=np.bool)
else:
mask = mask != 0
if estimate.shape != mask.shape:
raise Exception(f'estimate and mask have to be same shape (expected {estimate.shape} == {mask.shape})')
return estimate, target, mask
estimate = torch.squeeze(estimate[:, 0, :, :])
target = torch.squeeze(target[:, 0, :, :])
estimate, target, mask = _process_inputs(estimate, target, mask)
mask = mask.cpu().detach().numpy()
estimate = estimate.cpu().detach().numpy()
target = target.cpu().detach().numpy()
diff = np.abs(estimate[mask] - target[mask])
m = (diff > threshold).sum() / mask.sum()
return m
def ensure_dir(path):
if not os.path.exists(path):
os.makedirs(path, exist_ok=True)
@ -134,12 +160,12 @@ def sequence_loss(flow_preds, flow_gt, valid, gamma=0.8, test=False):
return flow_loss
class CREStereoLightning(LightningModule):
def __init__(self, args, logger, pattern_path, data_path):
def __init__(self, args, logger=None, pattern_path='', data_path=''):
super().__init__()
self.batch_size = args.batch_size
self.wandb_logger = logger
self.data_type = 'blender' if 'blender' in data_path else 'ctd'
self.lr = args.base_lr
print(f'lr = {self.lr}')
self.T_max = args.t_max if args.t_max else None
@ -149,13 +175,25 @@ class CREStereoLightning(LightningModule):
self.model = Model(
max_disp=args.max_disp, mixed_precision=args.mixed_precision, test_mode=False
)
# so I can access it in adjust learn rate more easily
self.n_total_epoch = args.n_total_epoch
self.base_lr = args.base_lr
self.automatic_optimization = False
def train_dataloader(self):
dataset = BlenderDataset(
root=self.data_path,
pattern_path=self.pattern_path,
use_lightning=True,
)
if self.data_type == 'blender':
dataset = BlenderDataset(
root=self.data_path,
pattern_path=self.pattern_path,
use_lightning=True,
)
elif self.data_type == 'ctd':
dataset = CTDDataset(
root=self.data_path,
pattern_path=self.pattern_path,
use_lightning=True,
)
dataloader = DataLoader(
dataset,
self.batch_size,
@ -169,12 +207,20 @@ class CREStereoLightning(LightningModule):
return dataloader
def val_dataloader(self):
test_dataset = BlenderDataset(
root=self.data_path,
pattern_path=self.pattern_path,
test_set=True,
use_lightning=True,
)
if self.data_type == 'blender':
test_dataset = BlenderDataset(
root=self.data_path,
pattern_path=self.pattern_path,
test_set=True,
use_lightning=True,
)
elif self.data_type == 'ctd':
test_dataset = CTDDataset(
root=self.data_path,
pattern_path=self.pattern_path,
test_set=True,
use_lightning=True,
)
test_dataloader = DataLoader(
test_dataset,
@ -190,12 +236,20 @@ class CREStereoLightning(LightningModule):
def test_dataloader(self):
# TODO change this to use IRL data?
test_dataset = BlenderDataset(
root=self.data_path,
pattern_path=self.pattern_path,
test_set=True,
use_lightning=True,
)
if self.data_type == 'blender':
test_dataset = CTDDataset(
root=self.data_path,
pattern_path=self.pattern_path,
test_set=True,
use_lightning=True,
)
elif self.data_type == 'ctd':
test_dataset = BlenderDataset(
root=self.data_path,
pattern_path=self.pattern_path,
test_set=True,
use_lightning=True,
)
test_dataloader = DataLoader(
test_dataset,
@ -232,6 +286,10 @@ class CREStereoLightning(LightningModule):
image_log['key'] = 'debug_train'
self.wandb_logger.log_image(**image_log)
self.log("train_loss", loss)
# update learn rate every N epochs
if self.trainer.is_last_batch:
self.adjust_learning_rate()
return loss
def validation_step(self, batch, batch_idx):
@ -243,6 +301,11 @@ class CREStereoLightning(LightningModule):
flow_predictions, gt_flow, valid_mask, gamma=0.8
)
self.log("val_loss", val_loss)
of = {}
for threshold in [0.1, 0.5, 1, 2, 5]:
of[threshold] = outlier_fraction(flow_predictions[0], gt_flow, valid_mask, threshold)
self.log("outlier_fraction", of)
# print(', '.join(f'of{thr}={val}' for thr, val in of.items()))
if batch_idx % 8 == 0:
self.wandb_logger.log_image(**log_images(left, right, flow_predictions, gt_disp))
@ -256,6 +319,9 @@ class CREStereoLightning(LightningModule):
)
self.log("test_loss", test_loss)
self.wandb_logger.log_image(**log_images(left, right, flow_predictions, gt_disp))
def predict_step(self, batch, batch_idx, dataloader_idx=0):
return self(batch)
def configure_optimizers(self):
optimizer = optim.Adam(self.model.parameters(), lr=self.lr, betas=(0.9, 0.999))
@ -268,35 +334,70 @@ class CREStereoLightning(LightningModule):
'name': 'CosineAnnealingLRScheduler',
}
return [optimizer], [lr_scheduler]
def adjust_learning_rate(self):
optimizer = self.optimizers().optimizer
epoch = self.trainer.current_epoch+1
warm_up = 0.02
const_range = 0.6
min_lr_rate = 0.05
if epoch <= self.n_total_epoch * warm_up:
lr = (1 - min_lr_rate) * self.base_lr / (
self.n_total_epoch * warm_up
) * epoch + min_lr_rate * self.base_lr
elif self.n_total_epoch * warm_up < epoch <= self.n_total_epoch * const_range:
lr = self.base_lr
else:
lr = (min_lr_rate - 1) * self.base_lr / (
(1 - const_range) * self.n_total_epoch
) * epoch + (1 - min_lr_rate * const_range) / (1 - const_range) * self.base_lr
for param_group in optimizer.param_groups:
param_group['lr'] = lr
if __name__ == "__main__":
wandb_logger = WandbLogger(project="crestereo-lightning", log_model=True)
# train configuration
args = parse_yaml("cfgs/train.yaml")
pattern_path = '/home/nils/miniprojekt/kinect_syn_ref.png'
run = wandb.init(project="crestereo-lightning", config=args._asdict(), tags=['new_scheduler', 'default_lr', f'{"" if args.pattern_attention else "no-"}pattern-attention'], notes='')
run.config.update(args._asdict())
wandb_logger.experiment.config.update(args._asdict())
config = wandb.config
wandb_logger = WandbLogger(project="crestereo-lightning", id=run.id, log_model=True)
# wandb_logger = WandbLogger(project="crestereo-lightning", log_model='all')
# wandb_logger.experiment.config.update(args._asdict())
if 'blender' in config.training_data_path:
# this was used for our blender renders
pattern_path = '/home/nils/miniprojekt/kinect_syn_ref.png'
if 'ctd' in config.training_data_path:
# this one is used (i hope) for ctd
pattern_path = '/home/nils/kinect_from_settings.png'
devices = min(config.nr_gpus, torch.cuda.device_count())
if devices != config.nr_gpus:
print(f'Using less devices than expected! ({devices} / {config.nr_gpus})')
model = CREStereoLightning(
# args,
config,
wandb_logger,
pattern_path,
args.training_data_path,
config.training_data_path,
# lr=0.00017378008287493763, # found with auto_lr_find=True
)
# NOTE turn this down once it's working, this might use too much space
# wandb_logger.watch(model, log_graph=False) #, log='all')
model_checkpoint = ModelCheckpoint(
monitor="val_loss",
mode="min",
save_top_k=2,
save_last=True,
)
trainer = Trainer(
accelerator='gpu',
devices=args.nr_gpus,
max_epochs=args.n_total_epoch,
devices=devices,
max_epochs=config.n_total_epoch,
callbacks=[
EarlyStopping(
monitor="val_loss",
@ -304,25 +405,20 @@ if __name__ == "__main__":
patience=16,
),
LearningRateMonitor(),
ModelCheckpoint(
monitor="val_loss",
mode="min",
save_top_k=2,
save_last=True,
)
model_checkpoint,
],
strategy=DDPSpawnStrategy(find_unused_parameters=False),
# auto_scale_batch_size='binsearch',
# auto_lr_find=True,
accumulate_grad_batches=4,
# accumulate_grad_batches=4, # needed to disable for manual optimization
deterministic=True,
check_val_every_n_epoch=1,
limit_val_batches=64,
limit_test_batches=256,
logger=wandb_logger,
default_root_dir=args.log_dir_lightning,
default_root_dir=config.log_dir_lightning,
)
# trainer.tune(model)
trainer.fit(model)
trainer.validate()
# trainer.validate(chkpt_path=model_checkpoint.best_model_path)

Loading…
Cancel
Save