def load_model(path, model):
for file in os.listdir(path):
if model in file:
print(f"Loading the following model for sampling: {file}")
try:
ckpt = torchlib.load_checkpoint(os.path.join(path, file))
if model == "(1)":
G_normal.load_state_dict(ckpt['G'])
return G_normal
elif model == "spectral":
G_spectral.load_state_dict(ckpt['G'])
return G_spectral
except Exception as e:
print(str(e))
return None
@torch.no_grad()
def sample(z):
G.eval()
return G(z)
# ==============================================================================
# = run =
# ==============================================================================
# load checkpoint if exists
ckpt_dir = py.join(output_dir, 'checkpoints')
py.mkdir(ckpt_dir)
try:
ckpt = torchlib.load_checkpoint(ckpt_dir)
ep, it_d, it_g = ckpt['ep'], ckpt['it_d'], ckpt['it_g']
D.load_state_dict(ckpt['D'])
G.load_state_dict(ckpt['G'])
D_optimizer.load_state_dict(ckpt['D_optimizer'])
G_optimizer.load_state_dict(ckpt['G_optimizer'])
except:
ep, it_d, it_g = 0, 0, 0
# sample
sample_dir = py.join(output_dir, 'samples_training')
py.mkdir(sample_dir)
# main loop
writer = tensorboardX.SummaryWriter(py.join(output_dir, 'summaries'))
z = torch.randn(100, args.z_dim, 1, 1).to(device) # a fixed noise for sampling
@torch.no_grad()
def sample(z):
G.eval()
return G(z)
# ==============================================================================
# = run =
# ==============================================================================
# load checkpoint if exists
ckpt_dir = py.join(output_dir, 'checkpoints')
py.mkdir(ckpt_dir)
try:
ckpt = torchlib.load_checkpoint(py.join(ckpt_dir, 'Last.ckpt'))
ep, it_d, it_g = ckpt['ep'], ckpt['it_d'], ckpt['it_g']
D.load_state_dict(ckpt['D'])
G.load_state_dict(ckpt['G'])
D_optimizer.load_state_dict(ckpt['D_optimizer'])
G_optimizer.load_state_dict(ckpt['G_optimizer'])
print('loading was successful. Starting at epoch: ', ep)
except Exception as e:
print(e)
ep, it_d, it_g = 0, 0, 0
# sample
sample_dir = py.join(output_dir, 'samples_training')
py.mkdir(sample_dir)
# main loop