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 seed: 0
mixed_precision: false mixed_precision: true
# base_lr: 4.0e-4 base_lr: 4.0e-4
base_lr: 0.001 # base_lr: 0.00001
t_max: 161 t_max: 16100
nr_gpus: 3 nr_gpus: 3
batch_size: 2 batch_size: 3
n_total_epoch: 300 n_total_epoch: 100
minibatch_per_epoch: 500 minibatch_per_epoch: 500
loadmodel: ~ loadmodel: ~
@ -17,11 +17,18 @@ model_save_freq_epoch: 1
max_disp: 256 max_disp: 256
image_width: 640 image_width: 640
image_height: 480 image_height: 480
# dataset: "blender"
# training_data_path: "./stereo_trainset/crestereo" # 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/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" log_level: "logging.INFO"

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

@ -94,6 +94,32 @@ def format_time(elapse):
return "{:02d}:{:02d}:{:02d}".format(hour, minute, seconds) 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): def ensure_dir(path):
if not os.path.exists(path): if not os.path.exists(path):
os.makedirs(path, exist_ok=True) 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 return flow_loss
class CREStereoLightning(LightningModule): class CREStereoLightning(LightningModule):
def __init__(self, args, logger, pattern_path, data_path): def __init__(self, args, logger=None, pattern_path='', data_path=''):
super().__init__() super().__init__()
self.batch_size = args.batch_size self.batch_size = args.batch_size
self.wandb_logger = logger self.wandb_logger = logger
self.data_type = 'blender' if 'blender' in data_path else 'ctd'
self.lr = args.base_lr self.lr = args.base_lr
print(f'lr = {self.lr}') print(f'lr = {self.lr}')
self.T_max = args.t_max if args.t_max else None self.T_max = args.t_max if args.t_max else None
@ -149,13 +175,25 @@ class CREStereoLightning(LightningModule):
self.model = Model( self.model = Model(
max_disp=args.max_disp, mixed_precision=args.mixed_precision, test_mode=False 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): def train_dataloader(self):
dataset = BlenderDataset( if self.data_type == 'blender':
root=self.data_path, dataset = BlenderDataset(
pattern_path=self.pattern_path, root=self.data_path,
use_lightning=True, 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( dataloader = DataLoader(
dataset, dataset,
self.batch_size, self.batch_size,
@ -169,12 +207,20 @@ class CREStereoLightning(LightningModule):
return dataloader return dataloader
def val_dataloader(self): def val_dataloader(self):
test_dataset = BlenderDataset( if self.data_type == 'blender':
root=self.data_path, test_dataset = BlenderDataset(
pattern_path=self.pattern_path, root=self.data_path,
test_set=True, pattern_path=self.pattern_path,
use_lightning=True, 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_dataloader = DataLoader(
test_dataset, test_dataset,
@ -190,12 +236,20 @@ class CREStereoLightning(LightningModule):
def test_dataloader(self): def test_dataloader(self):
# TODO change this to use IRL data? # TODO change this to use IRL data?
test_dataset = BlenderDataset( if self.data_type == 'blender':
root=self.data_path, test_dataset = CTDDataset(
pattern_path=self.pattern_path, root=self.data_path,
test_set=True, pattern_path=self.pattern_path,
use_lightning=True, 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_dataloader = DataLoader(
test_dataset, test_dataset,
@ -232,6 +286,10 @@ class CREStereoLightning(LightningModule):
image_log['key'] = 'debug_train' image_log['key'] = 'debug_train'
self.wandb_logger.log_image(**image_log) self.wandb_logger.log_image(**image_log)
self.log("train_loss", loss) self.log("train_loss", loss)
# update learn rate every N epochs
if self.trainer.is_last_batch:
self.adjust_learning_rate()
return loss return loss
def validation_step(self, batch, batch_idx): def validation_step(self, batch, batch_idx):
@ -243,6 +301,11 @@ class CREStereoLightning(LightningModule):
flow_predictions, gt_flow, valid_mask, gamma=0.8 flow_predictions, gt_flow, valid_mask, gamma=0.8
) )
self.log("val_loss", val_loss) 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: if batch_idx % 8 == 0:
self.wandb_logger.log_image(**log_images(left, right, flow_predictions, gt_disp)) self.wandb_logger.log_image(**log_images(left, right, flow_predictions, gt_disp))
@ -257,6 +320,9 @@ class CREStereoLightning(LightningModule):
self.log("test_loss", test_loss) self.log("test_loss", test_loss)
self.wandb_logger.log_image(**log_images(left, right, flow_predictions, gt_disp)) 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): def configure_optimizers(self):
optimizer = optim.Adam(self.model.parameters(), lr=self.lr, betas=(0.9, 0.999)) optimizer = optim.Adam(self.model.parameters(), lr=self.lr, betas=(0.9, 0.999))
print('len(self.train_dataloader)', len(self.train_dataloader())) print('len(self.train_dataloader)', len(self.train_dataloader()))
@ -269,34 +335,69 @@ class CREStereoLightning(LightningModule):
} }
return [optimizer], [lr_scheduler] 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__": if __name__ == "__main__":
wandb_logger = WandbLogger(project="crestereo-lightning", log_model=True)
# train configuration # train configuration
args = parse_yaml("cfgs/train.yaml") args = parse_yaml("cfgs/train.yaml")
pattern_path = '/home/nils/miniprojekt/kinect_syn_ref.png' wandb_logger.experiment.config.update(args._asdict())
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())
config = wandb.config config = wandb.config
wandb_logger = WandbLogger(project="crestereo-lightning", id=run.id, log_model=True) if 'blender' in config.training_data_path:
# wandb_logger = WandbLogger(project="crestereo-lightning", log_model='all') # this was used for our blender renders
# wandb_logger.experiment.config.update(args._asdict()) 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( model = CREStereoLightning(
# args, # args,
config, config,
wandb_logger, wandb_logger,
pattern_path, pattern_path,
args.training_data_path, config.training_data_path,
# lr=0.00017378008287493763, # found with auto_lr_find=True # lr=0.00017378008287493763, # found with auto_lr_find=True
) )
# NOTE turn this down once it's working, this might use too much space # NOTE turn this down once it's working, this might use too much space
# wandb_logger.watch(model, log_graph=False) #, log='all') # 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( trainer = Trainer(
accelerator='gpu', accelerator='gpu',
devices=args.nr_gpus, devices=devices,
max_epochs=args.n_total_epoch, max_epochs=config.n_total_epoch,
callbacks=[ callbacks=[
EarlyStopping( EarlyStopping(
monitor="val_loss", monitor="val_loss",
@ -304,25 +405,20 @@ if __name__ == "__main__":
patience=16, patience=16,
), ),
LearningRateMonitor(), LearningRateMonitor(),
ModelCheckpoint( model_checkpoint,
monitor="val_loss",
mode="min",
save_top_k=2,
save_last=True,
)
], ],
strategy=DDPSpawnStrategy(find_unused_parameters=False), strategy=DDPSpawnStrategy(find_unused_parameters=False),
# auto_scale_batch_size='binsearch', # auto_scale_batch_size='binsearch',
# auto_lr_find=True, # auto_lr_find=True,
accumulate_grad_batches=4, # accumulate_grad_batches=4, # needed to disable for manual optimization
deterministic=True, deterministic=True,
check_val_every_n_epoch=1, check_val_every_n_epoch=1,
limit_val_batches=64, limit_val_batches=64,
limit_test_batches=256, limit_test_batches=256,
logger=wandb_logger, logger=wandb_logger,
default_root_dir=args.log_dir_lightning, default_root_dir=config.log_dir_lightning,
) )
# trainer.tune(model) # trainer.tune(model)
trainer.fit(model) trainer.fit(model)
trainer.validate() # trainer.validate(chkpt_path=model_checkpoint.best_model_path)

Loading…
Cancel
Save