382 lines
11 KiB
Python
382 lines
11 KiB
Python
import argparse
|
|
import inspect
|
|
import os
|
|
import random
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
import paddle
|
|
import paddle.nn as nn
|
|
import paddle.optimizer as optim
|
|
from loguru import logger
|
|
from paddle.io import DataLoader
|
|
from paddle.nn import functional as F
|
|
from paddle_msssim import ms_ssim, ssim
|
|
|
|
from doc3d_dataset import Doc3dDataset
|
|
from GeoTr import GeoTr
|
|
from utils import to_image
|
|
|
|
FILE = Path(__file__).resolve()
|
|
ROOT = FILE.parents[0]
|
|
RANK = int(os.getenv("RANK", -1))
|
|
|
|
|
|
def init_seeds(seed=0, deterministic=False):
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
paddle.seed(seed)
|
|
if deterministic:
|
|
os.environ["FLAGS_cudnn_deterministic"] = "1"
|
|
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
|
|
os.environ["PYTHONHASHSEED"] = str(seed)
|
|
|
|
|
|
def colorstr(*input):
|
|
# Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code,
|
|
# i.e. colorstr('blue', 'hello world')
|
|
|
|
*args, string = (
|
|
input if len(input) > 1 else ("blue", "bold", input[0])
|
|
) # color arguments, string
|
|
|
|
colors = {
|
|
"black": "\033[30m", # basic colors
|
|
"red": "\033[31m",
|
|
"green": "\033[32m",
|
|
"yellow": "\033[33m",
|
|
"blue": "\033[34m",
|
|
"magenta": "\033[35m",
|
|
"cyan": "\033[36m",
|
|
"white": "\033[37m",
|
|
"bright_black": "\033[90m", # bright colors
|
|
"bright_red": "\033[91m",
|
|
"bright_green": "\033[92m",
|
|
"bright_yellow": "\033[93m",
|
|
"bright_blue": "\033[94m",
|
|
"bright_magenta": "\033[95m",
|
|
"bright_cyan": "\033[96m",
|
|
"bright_white": "\033[97m",
|
|
"end": "\033[0m", # misc
|
|
"bold": "\033[1m",
|
|
"underline": "\033[4m",
|
|
}
|
|
|
|
return "".join(colors[x] for x in args) + f"{string}" + colors["end"]
|
|
|
|
|
|
def print_args(args: Optional[dict] = None, show_file=True, show_func=False):
|
|
# Print function arguments (optional args dict)
|
|
x = inspect.currentframe().f_back # previous frame
|
|
file, _, func, _, _ = inspect.getframeinfo(x)
|
|
if args is None: # get args automatically
|
|
args, _, _, frm = inspect.getargvalues(x)
|
|
args = {k: v for k, v in frm.items() if k in args}
|
|
try:
|
|
file = Path(file).resolve().relative_to(ROOT).with_suffix("")
|
|
except ValueError:
|
|
file = Path(file).stem
|
|
s = (f"{file}: " if show_file else "") + (f"{func}: " if show_func else "")
|
|
logger.info(colorstr(s) + ", ".join(f"{k}={v}" for k, v in args.items()))
|
|
|
|
|
|
def increment_path(path, exist_ok=False, sep="", mkdir=False):
|
|
# Increment file or directory path,
|
|
# i.e. runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ... etc.
|
|
path = Path(path) # os-agnostic
|
|
if path.exists() and not exist_ok:
|
|
path, suffix = (
|
|
(path.with_suffix(""), path.suffix) if path.is_file() else (path, "")
|
|
)
|
|
|
|
for n in range(2, 9999):
|
|
p = f"{path}{sep}{n}{suffix}" # increment path
|
|
if not os.path.exists(p):
|
|
break
|
|
path = Path(p)
|
|
|
|
if mkdir:
|
|
path.mkdir(parents=True, exist_ok=True) # make directory
|
|
|
|
return path
|
|
|
|
|
|
def train(args):
|
|
save_dir = Path(args.save_dir)
|
|
|
|
use_vdl = args.use_vdl
|
|
if use_vdl:
|
|
from visualdl import LogWriter
|
|
|
|
log_dir = save_dir / "vdl"
|
|
vdl_writer = LogWriter(str(log_dir))
|
|
|
|
# Directories
|
|
weights_dir = save_dir / "weights"
|
|
weights_dir.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
last = weights_dir / "last.ckpt"
|
|
best = weights_dir / "best.ckpt"
|
|
|
|
# Hyperparameters
|
|
|
|
# Config
|
|
init_seeds(args.seed)
|
|
|
|
# Train loader
|
|
train_dataset = Doc3dDataset(
|
|
args.data_root,
|
|
split="train",
|
|
is_augment=True,
|
|
image_size=args.img_size,
|
|
)
|
|
|
|
train_loader = DataLoader(
|
|
train_dataset,
|
|
batch_size=args.batch_size,
|
|
shuffle=True,
|
|
num_workers=args.workers,
|
|
)
|
|
|
|
# Validation loader
|
|
val_dataset = Doc3dDataset(
|
|
args.data_root,
|
|
split="val",
|
|
is_augment=False,
|
|
image_size=args.img_size,
|
|
)
|
|
|
|
val_loader = DataLoader(
|
|
val_dataset, batch_size=args.batch_size, num_workers=args.workers
|
|
)
|
|
|
|
# Model
|
|
model = GeoTr()
|
|
|
|
if use_vdl:
|
|
vdl_writer.add_graph(
|
|
model,
|
|
input_spec=[
|
|
paddle.static.InputSpec([1, 3, args.img_size, args.img_size], "float32")
|
|
],
|
|
)
|
|
|
|
# Data Parallel Mode
|
|
if RANK == -1 and paddle.device.cuda.device_count() > 1:
|
|
model = paddle.DataParallel(model)
|
|
|
|
# Scheduler
|
|
scheduler = optim.lr.OneCycleLR(
|
|
max_learning_rate=args.lr,
|
|
total_steps=args.epochs * len(train_loader),
|
|
phase_pct=0.1,
|
|
end_learning_rate=args.lr / 2.5e5,
|
|
)
|
|
|
|
# Optimizer
|
|
optimizer = optim.AdamW(
|
|
learning_rate=scheduler,
|
|
parameters=model.parameters(),
|
|
)
|
|
|
|
# loss function
|
|
l1_loss_fn = nn.L1Loss()
|
|
mse_loss_fn = nn.MSELoss()
|
|
|
|
# Resume
|
|
best_fitness, start_epoch = 0.0, 0
|
|
if args.resume:
|
|
ckpt = paddle.load(args.resume)
|
|
model.set_state_dict(ckpt["model"])
|
|
optimizer.set_state_dict(ckpt["optimizer"])
|
|
scheduler.set_state_dict(ckpt["scheduler"])
|
|
best_fitness = ckpt["best_fitness"]
|
|
start_epoch = ckpt["epoch"] + 1
|
|
|
|
# Train
|
|
for epoch in range(start_epoch, args.epochs):
|
|
model.train()
|
|
|
|
for i, (img, target) in enumerate(train_loader):
|
|
img = paddle.to_tensor(img) # NCHW
|
|
target = paddle.to_tensor(target) # NHWC
|
|
|
|
pred = model(img) # NCHW
|
|
pred_nhwc = pred.transpose([0, 2, 3, 1])
|
|
|
|
loss = l1_loss_fn(pred_nhwc, target)
|
|
mse_loss = mse_loss_fn(pred_nhwc, target)
|
|
|
|
if use_vdl:
|
|
vdl_writer.add_scalar(
|
|
"Train/L1 Loss", float(loss), epoch * len(train_loader) + i
|
|
)
|
|
vdl_writer.add_scalar(
|
|
"Train/MSE Loss", float(mse_loss), epoch * len(train_loader) + i
|
|
)
|
|
vdl_writer.add_scalar(
|
|
"Train/Learning Rate",
|
|
float(scheduler.get_lr()),
|
|
epoch * len(train_loader) + i,
|
|
)
|
|
|
|
loss.backward()
|
|
optimizer.step()
|
|
scheduler.step()
|
|
optimizer.clear_grad()
|
|
|
|
if i % 10 == 0:
|
|
logger.info(
|
|
f"[TRAIN MODE] Epoch: {epoch}, Iter: {i}, L1 Loss: {float(loss)}, "
|
|
f"MSE Loss: {float(mse_loss)}, LR: {float(scheduler.get_lr())}"
|
|
)
|
|
|
|
# Validation
|
|
model.eval()
|
|
|
|
with paddle.no_grad():
|
|
avg_ssim = paddle.zeros([])
|
|
avg_ms_ssim = paddle.zeros([])
|
|
avg_l1_loss = paddle.zeros([])
|
|
avg_mse_loss = paddle.zeros([])
|
|
|
|
for i, (img, target) in enumerate(val_loader):
|
|
img = paddle.to_tensor(img)
|
|
target = paddle.to_tensor(target)
|
|
|
|
pred = model(img)
|
|
pred_nhwc = pred.transpose([0, 2, 3, 1])
|
|
|
|
# predict image
|
|
out = F.grid_sample(img, (pred_nhwc / args.img_size - 0.5) * 2)
|
|
out_gt = F.grid_sample(img, (target / args.img_size - 0.5) * 2)
|
|
|
|
# calculate ssim
|
|
ssim_val = ssim(out, out_gt, data_range=1.0)
|
|
ms_ssim_val = ms_ssim(out, out_gt, data_range=1.0)
|
|
|
|
loss = l1_loss_fn(pred_nhwc, target)
|
|
mse_loss = mse_loss_fn(pred_nhwc, target)
|
|
|
|
# calculate fitness
|
|
avg_ssim += ssim_val
|
|
avg_ms_ssim += ms_ssim_val
|
|
avg_l1_loss += loss
|
|
avg_mse_loss += mse_loss
|
|
|
|
if i % 10 == 0:
|
|
logger.info(
|
|
f"[VAL MODE] Epoch: {epoch}, VAL Iter: {i}, "
|
|
f"L1 Loss: {float(loss)} MSE Loss: {float(mse_loss)}, "
|
|
f"MS-SSIM: {float(ms_ssim_val)}, SSIM: {float(ssim_val)}"
|
|
)
|
|
|
|
if use_vdl and i == 0:
|
|
img_0 = to_image(out[0])
|
|
img_gt_0 = to_image(out_gt[0])
|
|
vdl_writer.add_image("Val/Predicted Image No.0", img_0, epoch)
|
|
vdl_writer.add_image("Val/Target Image No.0", img_gt_0, epoch)
|
|
|
|
img_1 = to_image(out[1])
|
|
img_gt_1 = to_image(out_gt[1])
|
|
img_gt_1 = img_gt_1.astype("uint8")
|
|
vdl_writer.add_image("Val/Predicted Image No.1", img_1, epoch)
|
|
vdl_writer.add_image("Val/Target Image No.1", img_gt_1, epoch)
|
|
|
|
img_2 = to_image(out[2])
|
|
img_gt_2 = to_image(out_gt[2])
|
|
vdl_writer.add_image("Val/Predicted Image No.2", img_2, epoch)
|
|
vdl_writer.add_image("Val/Target Image No.2", img_gt_2, epoch)
|
|
|
|
avg_ssim /= len(val_loader)
|
|
avg_ms_ssim /= len(val_loader)
|
|
avg_l1_loss /= len(val_loader)
|
|
avg_mse_loss /= len(val_loader)
|
|
|
|
if use_vdl:
|
|
vdl_writer.add_scalar("Val/L1 Loss", float(loss), epoch)
|
|
vdl_writer.add_scalar("Val/MSE Loss", float(mse_loss), epoch)
|
|
vdl_writer.add_scalar("Val/SSIM", float(ssim_val), epoch)
|
|
vdl_writer.add_scalar("Val/MS-SSIM", float(ms_ssim_val), epoch)
|
|
|
|
# Save
|
|
ckpt = {
|
|
"model": model.state_dict(),
|
|
"optimizer": optimizer.state_dict(),
|
|
"scheduler": scheduler.state_dict(),
|
|
"best_fitness": best_fitness,
|
|
"epoch": epoch,
|
|
}
|
|
|
|
paddle.save(ckpt, str(last))
|
|
|
|
if best_fitness < avg_ssim:
|
|
best_fitness = avg_ssim
|
|
paddle.save(ckpt, str(best))
|
|
|
|
if use_vdl:
|
|
vdl_writer.close()
|
|
|
|
|
|
def main(args):
|
|
print_args(vars(args))
|
|
|
|
args.save_dir = str(
|
|
increment_path(Path(args.project) / args.name, exist_ok=args.exist_ok)
|
|
)
|
|
|
|
train(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Hyperparams")
|
|
parser.add_argument(
|
|
"--data-root",
|
|
nargs="?",
|
|
type=str,
|
|
default="~/datasets/doc3d",
|
|
help="The root path of the dataset",
|
|
)
|
|
parser.add_argument(
|
|
"--img-size",
|
|
nargs="?",
|
|
type=int,
|
|
default=288,
|
|
help="The size of the input image",
|
|
)
|
|
parser.add_argument(
|
|
"--epochs",
|
|
nargs="?",
|
|
type=int,
|
|
default=65,
|
|
help="The number of training epochs",
|
|
)
|
|
parser.add_argument(
|
|
"--batch-size", nargs="?", type=int, default=12, help="Batch Size"
|
|
)
|
|
parser.add_argument(
|
|
"--lr", nargs="?", type=float, default=1e-04, help="Learning Rate"
|
|
)
|
|
parser.add_argument(
|
|
"--resume",
|
|
nargs="?",
|
|
type=str,
|
|
default=None,
|
|
help="Path to previous saved model to restart from",
|
|
)
|
|
parser.add_argument("--workers", type=int, default=8, help="max dataloader workers")
|
|
parser.add_argument(
|
|
"--project", default=ROOT / "runs/train", help="save to project/name"
|
|
)
|
|
parser.add_argument("--name", default="exp", help="save to project/name")
|
|
parser.add_argument(
|
|
"--exist-ok",
|
|
action="store_true",
|
|
help="existing project/name ok, do not increment",
|
|
)
|
|
parser.add_argument("--seed", type=int, default=0, help="Global training seed")
|
|
parser.add_argument("--use-vdl", action="store_true", help="use VisualDL as logger")
|
|
args = parser.parse_args()
|
|
main(args)
|