#
import argparse
import os
import pdb # For debug
import torch
from data import get_data
from model import get_model, model_device, valid_epoch
if __name__ == "__main__":
"""Test model."""
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--checkpoint',
type=str,
default="output/{{ . }}.pth",
help="checkpoint file")
parser.add_argument('--bs', type=int, default=2, help="batch size")
args = parser.parse_args()
# get model
model = get_model(args.checkpoint)
device = model_device()
model = model.to(device)
print("Start testing ...")
test_dl = get_data(trainning=False, bs=args.bs)
valid_epoch(test_dl, model, device, tag='test')
model_load(model.net_D, args.checkpoint_d)
model.to(device)
lr_scheduler_G = optim.lr_scheduler.StepLR(
model.optimizer_G, step_size=100, gamma=0.1)
if model.use_D:
lr_scheduler_D = optim.lr_scheduler.StepLR(
model.optimizer_D, step_size=100, gamma=0.1)
# get data loader
train_dl, valid_dl = get_data(trainning=True, bs=args.bs)
for epoch in range(args.epochs):
print("Epoch {}/{}, learning rate: {} ...".format(epoch + 1,
args.epochs, lr_scheduler_G.get_last_lr()))
train_epoch(train_dl, model, device, tag='train')
valid_epoch(valid_dl, model, device, tag='valid')
lr_scheduler_G.step()
if model.use_D:
lr_scheduler_D.step()
if (epoch + 1) % 100 == 0 or (epoch == args.epochs - 1):
model_save(model.net_G, os.path.join(
args.outputdir, "ImageColor_G_{}.pth".format(epoch + 1)))
if (model.use_D):
model_save(model.net_G, os.path.join(
args.outputdir, "ImageColor_D_{}.pth".format(epoch + 1)))
