def test(cfg: Namespace) -> None:
assert cfg.checkpoint not in [None, ""]
assert cfg.device == "cpu" or (cfg.device == "cuda"
and T.cuda.is_available())
exp_dir = ROOT_EXP_DIR / cfg.exp_name
os.makedirs(exp_dir / "out", exist_ok=True)
cfg.to_file(exp_dir / "test_config.json")
logger.info(f"[exp dir={exp_dir}]")
model = CAE()
model.load_state_dict(T.load(cfg.checkpoint))
model.eval()
if cfg.device == "cuda":
model.cuda()
logger.info(f"[model={cfg.checkpoint}] on {cfg.device}")
dataloader = DataLoader(dataset=ImageFolder720p(cfg.dataset_path),
batch_size=1,
shuffle=cfg.shuffle)
logger.info(f"[dataset={cfg.dataset_path}]")
loss_criterion = nn.MSELoss()
for batch_idx, data in enumerate(dataloader, start=1):
img, patches, _ = data
if cfg.device == "cuda":
patches = patches.cuda()
if batch_idx % cfg.batch_every == 0:
pass
out = T.zeros(6, 10, 3, 128, 128)
avg_loss = 0
for i in range(6):
for j in range(10):
x = patches[:, :, i, j, :, :].cuda()
y = model(x)
out[i, j] = y.data
loss = loss_criterion(y, x)
avg_loss += (1 / 60) * loss.item()
logger.debug("[%5d/%5d] avg_loss: %f", batch_idx, len(dataloader),
avg_loss)
# save output
out = np.transpose(out, (0, 3, 1, 4, 2))
out = np.reshape(out, (768, 1280, 3))
out = np.transpose(out, (2, 0, 1))
y = T.cat((img[0], out), dim=2)
save_imgs(
imgs=y.unsqueeze(0),
to_size=(3, 768, 2 * 1280),
name=exp_dir / f"out/test_{batch_idx}.png",
)
def test(cfg: Namespace) -> None:
assert cfg.checkpoint not in [None, ""]
assert cfg.device == "cpu" or (cfg.device == "cuda"
and T.cuda.is_available())
exp_dir = ROOT_EXP_DIR / cfg.exp_name
os.makedirs(exp_dir / "out", exist_ok=True)
cfg.to_file(exp_dir / "test_config.json")
logger.info(f"[exp dir={exp_dir}]")
model = CAE()
model.load_state_dict(T.load(cfg.checkpoint))
model.eval()
if cfg.device == "cuda":
model.cuda()
logger.info(f"[model={cfg.checkpoint}] on {cfg.device}")
dataloader = DataLoader(dataset=ImageFolder720p(cfg.dataset_path),
batch_size=1,
shuffle=cfg.shuffle)
logger.info(f"[dataset={cfg.dataset_path}]")
loss_criterion = nn.MSELoss()
for batch_idx, data in enumerate(dataloader, start=2):
img, patches, _ = data
print('the patches shape is:', patches.shape)
# print(_)
# plt.imshow(patches[0,:,3,1,:,:].permute(1,2,0))
# # plt.imshow(patches[0].permute(1,2,0))
# plt.show()
if cfg.device == "cuda":
patches = patches.cuda()
if batch_idx % cfg.batch_every == 0:
pass
out = T.zeros(6, 10, 3, 128, 128)
avg_loss = 0
foo = []
for i in range(6):
for j in range(10):
x = patches[:, :, i, j, :, :].cuda()
print('the x shape is:', x.shape)
y = model(x)
print('hellyy', y.shape)
def test(cfg):
os.makedirs(f"./test/{cfg['exp_name']}", exist_ok=True)
model = CAE().cuda()
model.load_state_dict(torch.load(cfg['chkpt']))
model.eval()
logger.info("Loaded model from", cfg['chkpt'])
dataset = ImageFolder720p(cfg['dataset_path'])
dataloader = DataLoader(dataset, batch_size=1, shuffle=cfg['shuffle'])
logger.info(f"Done setup dataloader: {len(dataloader)}")
mse_loss = nn.MSELoss()
for bi, (img, patches, path) in enumerate(dataloader):
out = torch.zeros(6, 10, 3, 128, 128)
# enc = torch.zeros(6, 10, 16, 8, 8)
avg_loss = 0
for i in range(6):
for j in range(10):
x = Variable(patches[:, :, i, j, :, :]).cuda()
y = model(x)
# e = model.enc_x.data
# p = torch.tensor(np.random.permutation(e.reshape(-1, 1)).reshape(1, 16, 8, 8)).cuda()
# out[i, j] = model.decode(p).data
# enc[i, j] = model.enc_x.data
out[i, j] = y.data
loss = mse_loss(y, x)
avg_loss += (1 / 60) * loss.item()
logger.debug('[%5d/%5d] avg_loss: %f' % (bi, len(dataloader), avg_loss))
# save output
out = np.transpose(out, (0, 3, 1, 4, 2))
out = np.reshape(out, (768, 1280, 3))
out = np.transpose(out, (2, 0, 1))
y = torch.cat((img[0], out), dim=2)
save_imgs(imgs=y.unsqueeze(0), to_size=(3, 768, 2 * 1280), name=f"./test/{cfg['exp_name']}/test_{bi}.png")
def imgEncoding(name,patches,checkpoint,interFolder):
device = 'cuda' if torch.cuda.is_available() else 'cpu'
encoder = CAE()
encoder.load_state_dict(torch.load(checkpoint))
encoder.eval()
encoder.to(device)
for idx,patcher in enumerate(patches):
logging.info(f'Doing for {name[idx]}')
out = torch.zeros(6,10, 32, 32, 32)
for i in range(6):
for j in range(10):
x = patcher[None,:, i, j, :, :].to(device)
y = encoder(x)
out[i,j]=y.data
torch.save(out,os.path.join(interFolder,str(name[idx])+'.pt'))
logging.info('Images are encoded')
return encoder
def train(cfg):
os.makedirs(f"out/{cfg['exp_name']}", exist_ok=True)
os.makedirs(f"checkpoints/{cfg['exp_name']}", exist_ok=True)
# dump config for current experiment
with open(f"checkpoints/{cfg['exp_name']}/setup.cfg", "wt") as f:
for k, v in cfg.items():
f.write("%15s: %s\n" % (k, v))
model = CAE().cuda()
if cfg['load']:
model.load_state_dict(torch.load(cfg['chkpt']))
logger.info("Loaded model from", cfg['chkpt'])
model.train()
logger.info("Done setup model")
dataset = ImageFolder720p(cfg['dataset_path'])
dataloader = DataLoader(dataset,
batch_size=cfg['batch_size'],
shuffle=cfg['shuffle'],
num_workers=cfg['num_workers'])
logger.info(
f"Done setup dataloader: {len(dataloader)} batches of size {cfg['batch_size']}"
)
mse_loss = nn.MSELoss()
adam = torch.optim.Adam(model.parameters(),
lr=cfg['learning_rate'],
weight_decay=1e-5)
sgd = torch.optim.SGD(model.parameters(), lr=cfg['learning_rate'])
optimizer = adam
ra = 0
for ei in range(cfg['resume_epoch'], cfg['num_epochs']):
for bi, (img, patches, _) in enumerate(dataloader):
avg_loss = 0
for i in range(6):
for j in range(10):
x = Variable(patches[:, :, i, j, :, :]).cuda()
y = model(x)
loss = mse_loss(y, x)
avg_loss += (1 / 60) * loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
ra = avg_loss if bi == 0 else ra * bi / (bi + 1) + avg_loss / (bi +
1)
logger.debug('[%3d/%3d][%5d/%5d] avg_loss: %f, ra: %f' %
(ei + 1, cfg['num_epochs'], bi + 1, len(dataloader),
avg_loss, ra))
# save img
if bi % cfg['out_every'] == 0:
out = torch.zeros(6, 10, 3, 128, 128)
for i in range(6):
for j in range(10):
x = Variable(patches[0, :, i,
j, :, :].unsqueeze(0)).cuda()
out[i, j] = model(x).cpu().data
out = np.transpose(out, (0, 3, 1, 4, 2))
out = np.reshape(out, (768, 1280, 3))
out = np.transpose(out, (2, 0, 1))
y = torch.cat((img[0], out), dim=2).unsqueeze(0)
save_imgs(imgs=y,
to_size=(3, 768, 2 * 1280),
name=f"out/{cfg['exp_name']}/out_{ei}_{bi}.png")
# save model
if bi % cfg['save_every'] == cfg['save_every'] - 1:
torch.save(
model.state_dict(),
f"checkpoints/{cfg['exp_name']}/model_{ei}_{bi}.state")
# save final model
torch.save(model.state_dict(),
f"checkpoints/{cfg['exp_name']}/model_final.state")
import os
import yaml
import argparse
from pathlib import Path
import numpy as np
import torch as T
import torch.nn as nn
from torch.utils.data import DataLoader
from data_loader import ImageFolder720p
from utils import save_imgs
import matplotlib.pylab as plt
# from bagoftools.namespace import Namespace
# from bagoftools.logger import Logger
from models.cae_32x32x32_zero_pad_bin import CAE
model = CAE()
model.load_state_dict(T.load(r'../checkpoint\model_yt_small_final.state'))
encoded = T.load('something.pt', map_location=T.device('cpu'))
print(encoded.shape)
out = model.decode(encoded)
print(out.shape)
plt.imshow(out[0].detach().permute(1,2,0))
plt.show()
os.makedirs(exp_dir,exist_ok=True)
dataset = custom_single()
dataloader = DataLoader(
dataset=dataset,
batch_size=16,
shuffle=False,
num_workers=4,
)
model = CAE()
model.eval()
state_dict = torch.load(path)
model.load_state_dict(state_dict)
model = model.cuda()
# for batch_idx, data in enumerate(dataloader, start=1):
# patches, _ = data
# patches = patches.float().cuda()
# break
# out = T.zeros(33, 32, 3, 128, 128)
# all_patches = dataset.patches
# all_patches = all_patches.reshape(33,32,3,128,128)
# for i in range(33):
# for j in range(32):
# x = all_patches[i,j,...].unsqueeze(0).cuda().float()
# out[i, j] = model(x.float()).cpu().data