def run_me(rank=0,
Sgrid=[1],
Ngrid=[3],
nNgrid=1,
Ggrid=[25.],
nGgrid=1,
ngpus=3,
idx=0):
args = get_args()
args.name = "default"
cfg = get_cfg(args)
cfg.use_ddp = False
cfg.use_apex = False
gpuid = rank % ngpus # set gpuid
gpuid = 1
cfg.gpuid = gpuid
cfg.device = f"cuda:{gpuid}"
grid_idx = idx * (1 * ngpus) + rank
B_grid_idx = (grid_idx % 2)
N_grid_idx = (grid_idx // 2) % nNgrid
G_grid_idx = grid_idx // (nNgrid * 2) % nGgrid
S_grid_idx = grid_idx // (nGgrid * nNgrid * 2)
# -- force blind --
B_grid_idx = 0
# -- config settings --
cfg.use_collate = True
# cfg.dataset.download = False
# cfg.cls = cfg
cfg.S = Sgrid[S_grid_idx]
# cfg.dataset.name = "cifar10"
cfg.dataset.name = "voc"
cfg.blind = (B_grid_idx == 0)
cfg.blind = False
cfg.N = Ngrid[N_grid_idx]
# cfg.N = 30
cfg.dynamic.frames = cfg.N
cfg.noise_type = 'g'
cfg.noise_params['g']['stddev'] = Ggrid[G_grid_idx]
noise_level = Ggrid[G_grid_idx] # don't worry about
cfg.batch_size = 4
cfg.init_lr = 1e-3
cfg.unet_channels = 3
cfg.input_N = cfg.N - 1
cfg.epochs = 50
cfg.color_cat = True
cfg.log_interval = int(int(50000 / cfg.batch_size) / 100)
cfg.dynamic.bool = True
cfg.dynamic.ppf = 2
cfg.dynamic.random_eraser = False
cfg.dynamic.frame_size = 128
# cfg.dynamic.total_pixels = cfg.dynamic.ppf * cfg.N
cfg.dynamic.total_pixels = 6
cfg.load = False
cfg.input_noise = False
cfg.middle_frame_random_erase = False
cfg.input_noise_level = noise_level / 255.
cfg.input_with_middle_frame = False
if (cfg.blind == 0): # e.g. supervised is true
cfg.input_with_middle_frame = True
if cfg.input_with_middle_frame:
cfg.input_N = cfg.N
blind = "blind" if cfg.blind else "nonblind"
print(grid_idx, blind, cfg.N, Ggrid[G_grid_idx], gpuid)
# if blind == "nonblind": return
dynamic_str = "dynamic_input_noise" if cfg.input_noise else "dynamic"
if cfg.input_with_middle_frame: dynamic_str += "_wmf"
postfix = Path(
f"./{dynamic_str}/{cfg.dynamic.frame_size}_{cfg.dynamic.ppf}_{cfg.dynamic.total_pixels}/{cfg.S}/{blind}/{cfg.N}/{noise_level}/"
)
print(postfix, cfg.dynamic.total_pixels)
cfg.model_path = cfg.model_path / postfix
cfg.optim_path = cfg.optim_path / postfix
if not cfg.model_path.exists(): cfg.model_path.mkdir(parents=True)
if not cfg.optim_path.exists(): cfg.optim_path.mkdir(parents=True)
checkpoint = cfg.model_path / Path("checkpoint_{}.tar".format(cfg.epochs))
# if checkpoint.exists(): return
print("N: {} | Noise Level: {}".format(cfg.N,
cfg.noise_params['g']['stddev']))
torch.cuda.set_device(gpuid)
# load model
model = load_model(cfg)
optimizer = load_optimizer(cfg, model)
scheduler = load_scheduler(cfg, model, optimizer)
nparams = count_parameters(model)
print("Number of Trainable Parameters: {}".format(nparams))
# load data
# data,loader = load_dataset(cfg,'denoising')
data, loader = load_dataset(cfg, 'dynamic')
# data,loader = simulate_noisy_dataset(data,loaders,M,N)
# load criterion
criterion = nn.BCELoss()
if cfg.load:
fp = cfg.model_path / Path("checkpoint_30.tar")
model = load_model_fp(cfg, model, fp, 0)
cfg.current_epoch = 0
te_ave_psnr = {}
test_before = False
if test_before:
ave_psnr = test_loop(cfg, model, criterion, loader.te, -1)
print("PSNR before training {:2.3e}".format(ave_psnr))
return
if checkpoint.exists() and cfg.load:
model = load_model_fp(cfg, model, checkpoint, gpuid)
print("Loaded model.")
cfg.current_epoch = cfg.epochs
for epoch in range(cfg.current_epoch, cfg.epochs):
losses = train_loop(cfg, model, optimizer, criterion, loader.tr, epoch)
ave_psnr = test_loop(cfg, model, criterion, loader.te, epoch)
te_ave_psnr[epoch] = ave_psnr
cfg.current_epoch += 1
epochs, psnr = zip(*te_ave_psnr.items())
best_index = np.argmax(psnr)
best_epoch, best_psnr = epochs[best_index], psnr[best_index]
print(
f"Best Epoch {best_epoch} | Best PSNR {best_psnr} | N: {cfg.N} | Blind: {blind}"
)
root = Path(f"{settings.ROOT_PATH}/output/pib/{postfix}/")
# if cfg.blind: root = root / Path(f"./blind/")
# else: root = root / Path(f"./nonblind/")
fn = Path(f"results.csv")
if not root.exists(): root.mkdir(parents=True)
path = root / fn
with open(path, 'w') as f:
f.write("{:d},{:d},{:2.10e},{:d}\n".format(cfg.N, best_epoch,
best_psnr, nparams))
save_model(cfg, model, optimizer)
def run_me(tr_method,
rank=0,
Sgrid=[1],
Ngrid=[3],
nNgrid=1,
Ggrid=[75.],
nGgrid=1,
ngpus=3,
idx=0):
args = get_args()
args.name = "default"
cfg = get_cfg(args)
cfg.use_ddp = False
cfg.use_apex = False
cfg.global_steps = 0
gpuid = rank % ngpus # set gpuid
gpuid = 1
cfg.gpuid = gpuid
cfg.device = f"cuda:{gpuid}"
grid_idx = idx * (1 * ngpus) + rank
B_grid_idx = (grid_idx % 2)
N_grid_idx = (grid_idx // 2) % nNgrid
G_grid_idx = grid_idx // (nNgrid * 2) % nGgrid
S_grid_idx = grid_idx // (nGgrid * nNgrid * 2)
# -- force blind --
B_grid_idx = 0
# -- config settings --
cfg.use_collate = True
# cfg.dataset.download = False
# cfg.cls = cfg
cfg.S = Sgrid[S_grid_idx]
# cfg.dataset.name = "cifar10"
cfg.dataset.name = "voc"
# cfg.dataset.name = "rebel2021"
cfg.blind = (B_grid_idx == 0)
cfg.blind = False
cfg.N = Ngrid[N_grid_idx]
cfg.N = 5
cfg.use_anscombe = True
# -- noise 2 simulate parameters --
cfg.sim_shuffleK = True
cfg.sim_method = "l2"
cfg.sim_K = 1
cfg.sim_patchsize = 9
# cfg.N = 30
cfg.num_workers = 2
cfg.dynamic.frames = cfg.N
# -- gaussian noise --
noise_type = "g"
cfg.noise_type = noise_type
cfg.noise_params['g']['stddev'] = Ggrid[G_grid_idx]
cfg.noise_params.ntype = cfg.noise_type
noise_params = cfg.noise_params['g']
noise_level = Ggrid[G_grid_idx]
noise_level_str = f"{int(noise_level)}"
# -- heteroskedastic gaussian noise --
# noise_type = "hg"
# cfg.noise_type = noise_type
# cfg.noise_params['hg']['read'] = Ggrid[G_grid_idx]
# cfg.noise_params['hg']['shot'] = 25.
# noise_params = cfg.noise_params['hg']
# cfg.noise_params.ntype = cfg.noise_type
# noise_level = Ggrid[G_grid_idx]
# noise_level_str = f"{int(noise_params['read']),int(noise_params['shot'])}"
# -- low-light noise --
# noise_type = "qis"
# cfg.noise_type = noise_type
# cfg.noise_params['qis']['alpha'] = 4.0
# cfg.noise_params['qis']['readout'] = 0.0
# cfg.noise_params['qis']['nbits'] = 3
# noise_params = cfg.noise_params['qis']
# cfg.noise_params.ntype = cfg.noise_type
# noise_level = noise_params['readout']
# noise_level_str = f"{int(noise_params['alpha']),int(noise_params['readout']),int(noise_params['nbits'])}"
# -- batch info --
cfg.batch_size = 10
cfg.init_lr = 1e-4
cfg.unet_channels = 3
cfg.input_N = cfg.N - 1
cfg.epochs = 50
cfg.color_cat = True
cfg.log_interval = int(int(50000 / cfg.batch_size) / 100)
cfg.dynamic.bool = True
cfg.dynamic.ppf = 0
cfg.dynamic.random_eraser = False
cfg.dynamic.frame_size = 128
cfg.dynamic.total_pixels = cfg.dynamic.ppf * cfg.N
cfg.dynamic.nframe = cfg.N
# cfg.dynamic.total_pixels = 6
cfg.load = False
print(
"This is to benchmark a new distribution class. We want to find the case when mL2 fails."
)
cfg.input_noise = False
cfg.input_noise_middle_only = False
cfg.input_with_middle_frame = True
cfg.middle_frame_random_erase = False
cfg.input_noise_level = noise_level / 255.
if (cfg.blind == 0): # e.g. supervised is true
cfg.input_with_middle_frame = True
if cfg.input_with_middle_frame:
cfg.input_N = cfg.N
blind = "blind" if cfg.blind else "nonblind"
print(grid_idx, blind, cfg.N, Ggrid[G_grid_idx], gpuid)
# if blind == "nonblind": return
dynamic_str = "dynamic_input_noise" if cfg.input_noise else "dynamic"
if cfg.input_noise_middle_only: dynamic_str += "_mo"
if cfg.input_with_middle_frame: dynamic_str += "_wmf"
postfix = Path(
f"./{dynamic_str}/{cfg.dynamic.frame_size}_{cfg.dynamic.ppf}_{cfg.dynamic.total_pixels}/{cfg.S}/{blind}/{cfg.N}/{noise_level}/"
)
print(postfix, cfg.dynamic.total_pixels)
cfg.model_path = cfg.model_path / postfix
cfg.optim_path = cfg.optim_path / postfix
if not cfg.model_path.exists(): cfg.model_path.mkdir(parents=True)
if not cfg.optim_path.exists(): cfg.optim_path.mkdir(parents=True)
checkpoint = cfg.model_path / Path("checkpoint_{}.tar".format(cfg.epochs))
# if checkpoint.exists(): return
print(
f"Sim Method: {cfg.sim_method} | Shuffle K {cfg.sim_shuffleK} | Sim K: {cfg.sim_K} | Patchsize: {cfg.sim_patchsize}"
)
print(f"Ascombe Transform: {cfg.use_anscombe}")
print("N: {} | Noise Level: {} | Noise Type: {}".format(
cfg.N, noise_level_str, noise_type))
print("PID: {}".format(os.getpid()))
torch.cuda.set_device(gpuid)
# load model
torch.manual_seed(cfg.seed)
model = load_model(cfg)
optimizer = load_optimizer(cfg, model)
scheduler = load_scheduler(cfg, model, optimizer)
nparams = count_parameters(model)
print("Number of Trainable Parameters: {}".format(nparams))
# load data
# data,loader = load_dataset(cfg,'denoising')
data, loader = load_dataset(cfg, 'dynamic')
# data,loader = simulate_noisy_dataset(data,loaders,M,N)
# load criterion
criterion = nn.BCELoss()
if cfg.load:
fp = cfg.model_path / Path("checkpoint_30.tar")
model = load_model_fp(cfg, model, fp, 0)
cfg.current_epoch = 0
te_ave_psnr = {}
test_before = False
if test_before:
ave_psnr = test_loop(cfg, model, criterion, loader.te, -1)
print("PSNR before training {:2.3e}".format(ave_psnr))
return
if checkpoint.exists() and cfg.load:
model = load_model_fp(cfg, model, checkpoint, gpuid)
print("Loaded model.")
cfg.current_epoch = cfg.epochs
# run_ot_v_displacement(cfg,criterion,loader.te)
for epoch in range(cfg.current_epoch, cfg.epochs):
if tr_method == "mse":
losses = train_loop_mse(cfg, model, optimizer, criterion,
loader.tr, epoch)
ave_psnr = test_loop_mse(cfg, model, criterion, loader.te, epoch)
elif tr_method == "n2n":
losses = train_loop_n2n(cfg, model, optimizer, criterion,
loader.tr, epoch)
ave_psnr = test_loop_n2n(cfg, model, criterion, loader.te, epoch)
else:
raise ValueError(f"Uknown training method [{tr_method}]")
# losses = train_loop(cfg,model,optimizer,criterion,loader.tr,epoch)
# ave_psnr = test_loop(cfg,model,criterion,loader.te,epoch)
te_ave_psnr[epoch] = ave_psnr
cfg.current_epoch += 1
epochs, psnr = zip(*te_ave_psnr.items())
best_index = np.argmax(psnr)
best_epoch, best_psnr = epochs[best_index], psnr[best_index]
print(
f"Best Epoch {best_epoch} | Best PSNR {best_psnr} | N: {cfg.N} | Blind: {blind}"
)
root = Path(f"{settings.ROOT_PATH}/output/n2n/{postfix}/")
# if cfg.blind: root = root / Path(f"./blind/")
# else: root = root / Path(f"./nonblind/")
fn = Path(f"results.csv")
if not root.exists(): root.mkdir(parents=True)
path = root / fn
with open(path, 'w') as f:
f.write("{:d},{:d},{:2.10e},{:d}\n".format(cfg.N, best_epoch,
best_psnr, nparams))
save_model(cfg, model, optimizer)
def save_denoising_model(cfg, model, optimizer):
save_model(cfg, model, None)
save_optim(cfg, optimizer)
def run_me(rank=0,
Sgrid=[50000],
Ngrid=[3],
nNgrid=1,
Ggrid=[25.],
nGgrid=1,
ngpus=3,
idx=0):
# def run_me(rank=1,Ngrid=1,Ggrid=1,nNgrid=1,ngpus=3,idx=1):
"""
PSNR 20 = (can equal) = AWGN @ 25
PSNR 25 = (can equal) = AWGN @ 14
PSNR 28 = (can equal) = AWGN @ 5
"""
args = get_args()
args.name = "default"
cfg = get_cfg(args)
cfg.use_ddp = False
cfg.use_apex = False
gpuid = 0
cfg.gpuid = gpuid
# gpuid = rank % ngpus # set gpuid
cfg.device = f"cuda:{gpuid}"
# -- experiment info --
cfg.exp_name = "sup_stn_standard"
cfg.desc = "Desc: sup stn model standard"
grid_idx = idx * (1 * ngpus) + rank
B_grid_idx = (grid_idx % 2)
N_grid_idx = (grid_idx // 2) % nNgrid
G_grid_idx = grid_idx // (nNgrid * 2) % nGgrid
S_grid_idx = grid_idx // (nGgrid * nNgrid * 2)
cfg.use_collate = True
# cfg.dataset.download = False
# cfg.cls = cfg
cfg.S = Sgrid[S_grid_idx]
# cfg.dataset.name = "cifar10"
cfg.dataset.name = "voc"
# cfg.blind = (B_grid_idx == 0)
cfg.blind = False
cfg.N = Ngrid[N_grid_idx]
cfg.N = 3 # 10
cfg.dynamic.frames = cfg.N
cfg.noise_type = 'g'
cfg.noise_params['g']['stddev'] = Ggrid[G_grid_idx]
noise_level = Ggrid[G_grid_idx]
cfg.batch_size = 4
cfg.init_lr = 1e-4
cfg.unet_channels = 3
cfg.input_N = cfg.N - 1
cfg.epochs = 100
cfg.log_interval = 50 # int(int(50000 / cfg.batch_size) / 100)
cfg.dynamic.bool = True
cfg.dynamic.ppf = 2
cfg.dynamic.frame_size = 256
cfg.dynamic.total_pixels = 2 * cfg.N
cfg.load = False
# -- input noise for learning --
cfg.input_noise = False
cfg.input_noise_middle_only = False
cfg.input_with_middle_frame = True
cfg.input_noise_level = noise_level / 255
if cfg.input_with_middle_frame:
cfg.input_N = cfg.N
blind = "blind" if cfg.blind else "nonblind"
print(grid_idx, blind, cfg.N, Ggrid[G_grid_idx], gpuid, cfg.input_noise,
cfg.input_with_middle_frame)
# if blind == "nonblind": return
dynamic_str = "dynamic_input_noise" if cfg.input_noise else "dynamic"
if cfg.input_noise_middle_only: dynamic_str += "_mo"
if cfg.input_with_middle_frame: dynamic_str += "_wmf"
postfix = Path(
f"./{dynamic_str}/{cfg.dynamic.frame_size}_{cfg.dynamic.ppf}_{cfg.dynamic.total_pixels}/{cfg.S}/{blind}/{cfg.N}/{noise_level}/"
)
print(postfix)
cfg.model_path = cfg.model_path / postfix
cfg.optim_path = cfg.optim_path / postfix
if not cfg.model_path.exists(): cfg.model_path.mkdir(parents=True)
if not cfg.optim_path.exists(): cfg.optim_path.mkdir(parents=True)
checkpoint = cfg.model_path / Path("checkpoint_{}.tar".format(cfg.epochs))
# if checkpoint.exists(): return
print("PID: {}".format(os.getpid()))
print("N: {} | Noise Level: {}".format(cfg.N,
cfg.noise_params['g']['stddev']))
torch.cuda.set_device(gpuid)
# -- load model --
# model,criterion = load_model_kpn(cfg)
model, criterion = load_model_stn(cfg)
optimizer = load_optimizer(cfg, model)
scheduler = load_scheduler(cfg, model, optimizer)
nparams = count_parameters(model)
print("Number of Trainable Parameters: {}".format(nparams))
# load data
# data,loader = load_dataset(cfg,'denoising')
data, loader = load_dataset(cfg, 'dynamic')
# data,loader = simulate_noisy_dataset(data,loaders,M,N)
if cfg.load:
fp = cfg.model_path / Path("checkpoint_30.tar")
model = load_model_fp(cfg, model, fp, 0)
cfg.current_epoch = 0
te_ave_psnr = {}
test_before = False
if test_before:
ave_psnr = test_loop(cfg, model, criterion, loader.te, -1)
print("PSNR before training {:2.3e}".format(ave_psnr))
return
if checkpoint.exists() and cfg.load:
model = load_model_fp(cfg, model, checkpoint, gpuid)
print("Loaded model.")
cfg.current_epoch = cfg.epochs
cfg.global_step = 0
use_record = False
record = init_record()
for epoch in range(cfg.current_epoch, cfg.epochs):
print(cfg.desc)
sys.stdout.flush()
losses, epoch_record = train_loop(cfg, model, optimizer, criterion,
loader.tr, epoch)
if use_record:
record = record.append(epoch_record)
write_record_file(cfg.current_epoch, postfix, record)
ave_psnr = test_loop(cfg, model, criterion, loader.te, epoch)
te_ave_psnr[epoch] = ave_psnr
cfg.current_epoch += 1
epochs, psnr = zip(*te_ave_psnr.items())
best_index = np.argmax(psnr)
best_epoch, best_psnr = epochs[best_index], psnr[best_index]
root = Path(f"{settings.ROOT_PATH}/output/n2n-stn/{postfix}/")
# if cfg.blind: root = root / Path(f"./blind/")
# else: root = root / Path(f"./nonblind/")
fn = Path(f"results.csv")
if not root.exists(): root.mkdir(parents=True)
path = root / fn
with open(path, 'w') as f:
f.write("{:d},{:d},{:2.10e},{:d}\n".format(cfg.N, best_epoch,
best_psnr, nparams))
save_model(cfg, model, optimizer)
def run_me(rank=0,Sgrid=[50000],Ngrid=[2],nNgrid=2,Ggrid=[25],nGgrid=1,ngpus=3,idx=0):
args = get_args()
args.name = "default"
cfg = get_cfg(args)
cfg.use_ddp = False
cfg.use_apex = False
# gpuid = rank % ngpus # set gpuid
gpuid = 0
cfg.device = f"cuda:{gpuid}"
grid_idx = idx*(1*ngpus)+rank
B_grid_idx = (grid_idx % 2)
N_grid_idx = ( grid_idx // 2 ) % nNgrid
G_grid_idx = grid_idx // (nNgrid * 2) % nGgrid
S_grid_idx = grid_idx // (nGgrid * nNgrid * 2)
cfg.use_collate = True
# cfg.dataset.download = False
# cfg.cls = cfg
cfg.S = Sgrid[S_grid_idx]
cfg.dataset.name = "cifar10"
# cfg.dataset.name = "voc"
cfg.blind = (B_grid_idx == 0)
cfg.blind = False
cfg.N = Ngrid[N_grid_idx]
cfg.dynamic.frames = cfg.N
cfg.noise_type = 'g'
cfg.noise_params['g']['stddev'] = Ggrid[G_grid_idx]
noise_level = Ggrid[G_grid_idx]
cfg.batch_size = 16
cfg.init_lr = 1e-3
cfg.unet_channels = 3
# if cfg.blind: cfg.input_N = cfg.N - 1
# else: cfg.input_N = cfg.N
cfg.input_N = cfg.N-1
cfg.epochs = 30
cfg.log_interval = int(int(50000 / cfg.batch_size) / 100)
cfg.dataset.load_residual = True
cfg.dynamic.bool = True
cfg.dynamic.ppf = 0
cfg.dynamic.frame_size = 128
cfg.dynamic.total_pixels = 0
cfg.load = False
blind = "blind" if cfg.blind else "nonblind"
print(grid_idx,blind,cfg.N,Ggrid[G_grid_idx],gpuid)
# if blind == "nonblind": return
postfix = get_postfix_str(cfg,blind,noise_level)
cfg.model_path = cfg.model_path / postfix
cfg.optim_path = cfg.optim_path / postfix
if not cfg.model_path.exists(): cfg.model_path.mkdir(parents=True)
if not cfg.optim_path.exists(): cfg.optim_path.mkdir(parents=True)
checkpoint = cfg.model_path / Path("checkpoint_{}.tar".format(cfg.epochs))
# if checkpoint.exists(): return
print("N: {} | Noise Level: {}".format(cfg.N,cfg.noise_params['g']['stddev']))
torch.cuda.set_device(gpuid)
# load model
model = DnCNN_Net(3)
optimizer = load_optimizer(cfg,model)
scheduler = load_scheduler(cfg,model,optimizer)
nparams = count_parameters(model)
print("Number of Trainable Parameters: {}".format(nparams))
# load data
data,loader = load_dataset(cfg,'denoising')
# data,loader = load_dataset(cfg,'dynamic')
# data,loader = simulate_noisy_dataset(data,loaders,M,N)
# load criterion
criterion = nn.BCELoss()
cfg.current_epoch = 0
te_ave_psnr = {}
test_before = False
if test_before:
ave_psnr = test_loop(cfg,model,criterion,loader.te,-1)
print("PSNR before training {:2.3e}".format(ave_psnr))
if checkpoint.exists() and cfg.load:
model = load_model_fp(cfg,model,checkpoint,gpuid)
print("Loaded model.")
cfg.current_epoch = cfg.epochs
for epoch in range(cfg.current_epoch,cfg.epochs):
losses = train_loop(cfg,model,optimizer,criterion,loader.tr,epoch)
ave_psnr = test_loop(cfg,model,criterion,loader.te,epoch)
te_ave_psnr[epoch] = ave_psnr
cfg.current_epoch += 1
epochs,psnr = zip(*te_ave_psnr.items())
best_index = np.argmax(psnr)
best_epoch,best_psnr = epochs[best_index],psnr[best_index]
root = Path(f"{settings.ROOT_PATH}/output/dncnn/{postfix}/")
# if cfg.blind: root = root / Path(f"./blind/")
# else: root = root / Path(f"./nonblind/")
fn = Path(f"results.csv")
if not root.exists(): root.mkdir(parents=True)
path = root / fn
with open(path,'w') as f:
f.write("{:d},{:d},{:2.10e},{:d}\n".format(cfg.N,best_epoch,best_psnr,nparams))
save_model(cfg, model, optimizer)
def save_simcl_model(cfg, model, optimizer):
print("Saving model and optim state")
save_model(cfg, model, None)
save_optim(cfg, optimizer)
def run_me():
args = get_args()
args.name = "default_attn_16_lowNoise"
args.dataset = "voc"
# args.num_workers = 0
cfg = get_cfg(args)
cfg.use_ddp = False
cfg.use_apex = False
gpuid = 1 # set gpuid
cfg.gpuid = gpuid
cfg.device = f"cuda:{gpuid}"
cfg.use_collate = True
cfg.dataset.download = False
cfg.cls = cfg
cfg.batch_size = 48
cfg.blind = False
cfg.N = 3
cfg.dynamic.frames = cfg.N
cfg.noise_type = 'g'
cfg.noise_params['g']['stddev'] = 25.
noise_level = 5.
cfg.init_lr = 1e-3
cfg.unet_channels = 3
cfg.input_N = cfg.N - 1
cfg.epochs = 1000
cfg.log_interval = int(int(50000 / cfg.batch_size) /
20) # NumOfBatches / Const = NumOfPrints
cfg.save_epoch = 50
# -- attn parameters --
cfg.patch_sizes = [48, 48]
cfg.d_model_attn = 3 # 512
cfg.dataset.name = "voc"
cfg.dataset.bw = False
cfg.dynamic.bool = True
cfg.dynamic.ppf = 2
cfg.dynamic.frame_size = 48
cfg.dynamic.total_pixels = 6
cfg.load = False
# left; 5 total pix; 10 in 8 out 16 imsize; bs 64; gpuid 1; 26 PSNR @ 30 and 190 epochs
# right; 20 total pix; 16 in 8 out 16 imsize; bs 12; gpuid 0; 21 PSNR @ 70 epochs
# 16_2_20 ; 27.59 psnr 16 in? 16 out? 16 imsize
# 10, 8 was also promising
# -- model type --
cfg.model_version = "v4"
# -- input noise for learning --
cfg.input_noise = True
cfg.input_noise_level = 25. / 255
cfg.input_with_middle_frame = True
if cfg.input_with_middle_frame:
cfg.input_N = cfg.N
# -- learning --
cfg.spoof_batch = False
cfg.spoof_batch_size = 1
torch.cuda.set_device(gpuid)
# cfg.batch_size = 256
cfg.init_lr = 5e-5
print(cfg.batch_size, cfg.init_lr)
# -- load model --
model = load_model(cfg)
# inputs = torch.rand(5,10,3,32,32).to(cfg.device)
# outputs = torch.rand(10,3,32,32).to(cfg.device)
# encode = model(inputs,outputs)
optimizer = load_optimizer(cfg, model)
# -- load data
data, loader = load_dataset(cfg, 'dynamic')
# data,loader = load_dataset(cfg,'denoising')
blind = "blind" if cfg.blind else "nonblind"
# print(grid_idx,blind,cfg.N,Ggrid[G_grid_idx],gpuid)
# if blind == "nonblind": return
dynamic_str = "dynamic_input_noise" if cfg.input_noise else "dynamic"
if cfg.input_with_middle_frame: dynamic_str += "_wmf"
postfix = Path(
f"./{cfg.model_version}/{dynamic_str}/{cfg.dynamic.frame_size}_{cfg.dynamic.ppf}_{cfg.dynamic.total_pixels}/{cfg.S}/{blind}/{cfg.N}/{noise_level}/"
)
print(f"model dir [{postfix}]")
cfg.model_path = cfg.model_path / postfix
cfg.optim_path = cfg.optim_path / postfix
if not cfg.model_path.exists(): cfg.model_path.mkdir(parents=True)
if not cfg.optim_path.exists(): cfg.optim_path.mkdir(parents=True)
checkpoint = cfg.model_path / Path("checkpoint_{}.tar".format(cfg.epochs))
print(checkpoint)
# if checkpoint.exists(): return
print("N: {} | Noise Level: {}".format(cfg.N,
cfg.noise_params['g']['stddev']))
torch.cuda.set_device(gpuid)
# load criterion
criterion = None
if cfg.load:
fp = cfg.model_path / Path("checkpoint_30.tar")
model = load_model_fp(cfg, model, fp, 0)
cfg.current_epoch = 0
te_ave_psnr = {}
test_before = False
# -- test before --
if test_before:
ave_psnr = test_loop(cfg, model, criterion, loader.te, -1)
print("PSNR before training {:2.3e}".format(ave_psnr))
return
# -- load trained model --
if checkpoint.exists() and cfg.load:
model = load_model_fp(cfg, model, checkpoint, gpuid)
print("Loaded model.")
cfg.current_epoch = cfg.epochs
# -- summary writer --
dsname = cfg.dataset.name
cfg.summary_log_dir = Path(
f"{settings.ROOT_PATH}/runs/attn/{dsname}/{cfg.exp_name}/{postfix}")
datetime_now = datetime.datetime.now().strftime("%b%d_%H-%M-%S")
writer_dir = cfg.summary_log_dir / Path(datetime_now)
writer = SummaryWriter(writer_dir)
print(f"Summary writer @ [{writer_dir}]")
# -- training loop --
for epoch in range(cfg.current_epoch, cfg.epochs):
losses = train_loop(cfg, model, optimizer, criterion, loader.tr, epoch)
writer.add_scalar("Loss/train", losses, epoch)
ave_psnr = test_loop(cfg, model, criterion, loader.te, epoch)
writer.add_scalar("Loss/test-ave-psnr", ave_psnr, epoch)
te_ave_psnr[epoch] = ave_psnr
cfg.current_epoch += 1
writer.flush()
if (epoch % cfg.save_epoch) == 0:
save_model(cfg, model, optimizer)
epochs, psnr = zip(*te_ave_psnr.items())
best_index = np.argmax(psnr)
best_epoch, best_psnr = epochs[best_index], psnr[best_index]
root = Path(f"{settings.ROOT_PATH}/output/attn/{postfix}/")
# if cfg.blind: root = root / Path(f"./blind/")
# else: root = root / Path(f"./nonblind/")
fn = Path(f"results.csv")
if not root.exists(): root.mkdir(parents=True)
path = root / fn
with open(path, 'w') as f:
f.write("{:d},{:d},{:2.10e},{:d}\n".format(cfg.N, best_epoch,
best_psnr, nparams))
save_model(cfg, model, optimizer)
def run_me(rank=0,Sgrid=[1],Ngrid=[3],nNgrid=1,Ggrid=[25.],nGgrid=1,ngpus=3,idx=0):
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
# algorithm hyperparameters
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
args = get_args()
args.name = "default"
cfg = get_cfg(args)
cfg.use_ddp = False
cfg.use_apex = False
gpuid = rank % ngpus # set gpuid
gpuid = 2
cfg.gpuid = gpuid
cfg.device = f"cuda:{gpuid}"
grid_idx = idx*(1*ngpus)+rank
B_grid_idx = (grid_idx % 2)
N_grid_idx = ( grid_idx // 2 ) % nNgrid
G_grid_idx = grid_idx // (nNgrid * 2) % nGgrid
S_grid_idx = grid_idx // (nGgrid * nNgrid * 2)
# -- experiment info --
cfg.exp_name = "unsup_byol_burst"
cfg.desc = "Desc: unsup byol burst"
# -- force blind --
B_grid_idx = 0
# -- config settings --
cfg.use_collate = True
# cfg.dataset.download = False
# cfg.cls = cfg
cfg.S = Sgrid[S_grid_idx]
# cfg.dataset.name = "cifar10"
cfg.dataset.name = "voc"
cfg.supervised = False
cfg.blind = (B_grid_idx == 0)
cfg.blind = ~cfg.supervised
cfg.N = Ngrid[N_grid_idx]
cfg.N = 3
# cfg.N = 30
cfg.dynamic.frames = cfg.N
cfg.noise_type = 'g'
cfg.noise_params['g']['stddev'] = Ggrid[G_grid_idx]
noise_level = Ggrid[G_grid_idx] # don't worry about
cfg.batch_size = 4
cfg.init_lr = 1e-4
cfg.unet_channels = 3
cfg.input_N = cfg.N-1
cfg.epochs = 100
cfg.color_cat = True
cfg.log_interval = 100 #int(int(50000 / cfg.batch_size) / 500)
cfg.dynamic.bool = True
cfg.dynamic.ppf = 2
cfg.dynamic.random_eraser = False
cfg.dynamic.frame_size = 64
# cfg.dynamic.total_pixels = cfg.dynamic.ppf * cfg.N
cfg.dynamic.total_pixels = 2*cfg.N
cfg.load = False
# -- attn params --
cfg.patch_sizes = [128,128]
cfg.d_model_attn = 3
cfg.input_noise = False
cfg.input_noise_middle_only = False
cfg.input_with_middle_frame = True
cfg.middle_frame_random_erase = False
cfg.input_noise_level = noise_level/255.
if (cfg.blind == 0): # e.g. supervised is true
cfg.input_with_middle_frame = True
if cfg.input_with_middle_frame:
cfg.input_N = cfg.N
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
# create strings from settings
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
blind = "blind" if cfg.blind else "nonblind"
print(grid_idx,blind,cfg.N,Ggrid[G_grid_idx],gpuid)
# if blind == "nonblind": return
dynamic_str = "dynamic_input_noise" if cfg.input_noise else "dynamic"
if cfg.input_noise_middle_only: dynamic_str += "_mo"
if cfg.input_with_middle_frame: dynamic_str += "_wmf"
postfix = Path(f"./modelBurst/{dynamic_str}/{cfg.dynamic.frame_size}_{cfg.dynamic.ppf}_{cfg.dynamic.total_pixels}/{cfg.S}/{blind}/{cfg.N}/{noise_level}/")
print(postfix,cfg.dynamic.total_pixels)
cfg.model_path = cfg.model_path / postfix
cfg.optim_path = cfg.optim_path / postfix
if not cfg.model_path.exists(): cfg.model_path.mkdir(parents=True)
if not cfg.optim_path.exists(): cfg.optim_path.mkdir(parents=True)
checkpoint = cfg.model_path / Path("checkpoint_{}.tar".format(cfg.epochs))
# if checkpoint.exists(): return
print("N: {} | Noise Level: {}".format(cfg.N,cfg.noise_params['g']['stddev']))
torch.cuda.set_device(gpuid)
# -=-=-=-=-=-=-=-=-=-
# load models
# -=-=-=-=-=-=-=-=-=-
# -- load model --
# model,criterion = load_burst_n2n_model(cfg)
model_online,criterion = load_burst_kpn_model(cfg)
optim_online = load_optimizer(cfg,model_online)
model_target,_ = load_burst_kpn_model(cfg)
optim_target = load_optimizer(cfg,model_target)
nparams = count_parameters(model_online)
print("Number of Trainable Parameters: {}".format(nparams))
print("PID: {}".format(os.getpid()))
# -=-=-=-=-=-=-=-=-=-
# load data
# -=-=-=-=-=-=-=-=-=-
# data,loader = load_dataset(cfg,'denoising')
data,loader = load_dataset(cfg,'dynamic')
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
# restore model from file
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
if cfg.load:
fp = cfg.model_path / Path("checkpoint_30.tar")
model = load_model_fp(cfg,model,fp,0)
cfg.current_epoch = 0
te_ave_psnr = {}
test_before = False
if test_before:
ave_psnr = test_loop_burstKPN(cfg,model,criterion,loader.te,-1)
print("PSNR before training {:2.3e}".format(ave_psnr))
return
if checkpoint.exists() and cfg.load:
model = load_model_fp(cfg,model,checkpoint,gpuid)
print("Loaded model.")
cfg.current_epoch = cfg.epochs
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
# train & test over epochs
# -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
record_losses = pd.DataFrame({'kpn':[],'ot':[],'psnr':[],'psnr_std':[]})
use_record = False
loss_type = "sup_r_ot"
for epoch in range(cfg.current_epoch,cfg.epochs):
print(cfg.desc)
sys.stdout.flush()
losses,record_losses = train_loop_burstKPN(cfg,model_target,model_online,optim_target,optim_online,criterion,loader.tr,epoch,record_losses)
if use_record:
write_record_losses_file(cfg.current_epoch,postfix,loss_type,record_losses)
ave_psnr,record_test = test_loop_burstKPN(cfg,model_online,criterion,loader.te,epoch)
if use_record:
write_record_test_file(cfg.current_epoch,postfix,loss_type,record_test)
te_ave_psnr[epoch] = ave_psnr
cfg.current_epoch += 1
epochs,psnr = zip(*te_ave_psnr.items())
best_index = np.argmax(psnr)
best_epoch,best_psnr = epochs[best_index],psnr[best_index]
print(f"Best Epoch {best_epoch} | Best PSNR {best_psnr} | N: {cfg.N} | Blind: {blind}")
root = Path(f"{settings.ROOT_PATH}/output/byol/{postfix}/")
# if cfg.blind: root = root / Path(f"./blind/")
# else: root = root / Path(f"./nonblind/")
fn = Path(f"results.csv")
if not root.exists(): root.mkdir(parents=True)
path = root / fn
with open(path,'w') as f:
f.write("{:d},{:d},{:2.10e},{:d}\n".format(cfg.N,best_epoch,best_psnr,nparams))
save_model(cfg, model, optimizer)