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
ROOT_EXP_DIR = Path(__file__).resolve().parents[1] / "experiments"
logger = Logger(__name__, colorize=True)
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":
from tensorboardX import SummaryWriter
from torch.autograd import Variable
from torch.utils.data import DataLoader
from data_loader import ImageFolder96p
from utils import get_config, get_args, dump_cfg
from utils import save_imgs
from bagoftools.logger import Logger
# models
sys.path.append(
os.path.join(os.path.dirname(os.path.realpath(__file__)), "../models"))
from cae_96x16x16_test import CAE
logger = Logger(name='train', colorize=True)
def prologue(cfg: Namespace, *varargs) -> SummaryWriter:
# sanity checks
assert cfg.device == "cpu" or (cfg.device == "cuda"
and T.cuda.is_available())
# dirs
base_dir = f"../experiments/{cfg.exp_name}"
os.makedirs(f"{base_dir}/out", exist_ok=True)
os.makedirs(f"{base_dir}/chkpt", exist_ok=True)
os.makedirs(f"{base_dir}/logs", exist_ok=True)
dump_cfg(f"{base_dir}/train_config.txt", vars(cfg))