def train(model, train_loader, valid_loader, batch_size, n_epochs, lr,
patience, step_size, device, model_type, save_path):
model_path = save_path + "best_model_" + model_type + ".pt"
model = model.to(device)
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = StepLR(optimizer, step_size=step_size, gamma=0.1)
train_losses = []
valid_losses = []
best_loss = math.inf
best_epoch = 0
k = 0
for epoch in range(n_epochs):
print("training.....")
train_loss = train_epoch(model, optimizer, epoch, train_loader, device,
model_type)
print("validation....")
valid_loss = validation(model, valid_loader, device, epoch, model_type)
train_losses.append(train_loss)
valid_losses.append(valid_loss)
print('Epoch {}: Train: avg. loss: {:.3f}'.format(
epoch + 1, train_loss))
print('Epoch {}: Valid: avg. loss: {:.3f}'.format(
epoch + 1, valid_loss))
if step_size != -1:
scheduler.step()
if valid_loss < best_loss:
best_loss = valid_loss
best_epoch = epoch + 1
print('Saving best model at epoch {}'.format(epoch + 1))
torch.save(model.state_dict(), model_path)
if model_type == "prediction":
gen_seq = generate_unconditional_seq(model_path,
700,
device,
bias=10.0,
style=None,
prime=False)
else:
gen_seq, phi = generate_conditional_sequence(
model_path,
"Hello world!",
device,
train_loader.dataset.char_to_id,
train_loader.dataset.idx_to_char,
bias=10.,
prime=False,
prime_seq=None,
real_text=None)
plt.imshow(phi, cmap='viridis', aspect='auto')
plt.colorbar()
plt.xlabel("time steps")
plt.yticks(np.arange(phi.shape[1]),
list("Hello world! "),
rotation='horizontal')
plt.margins(0.2)
plt.subplots_adjust(bottom=0.15)
plt.savefig(save_path + "heat_map" + str(best_epoch) + ".png")
plt.close()
# denormalize the generated offsets using train set mean and std
gen_seq = data_denormalization(Global.train_mean, Global.train_std,
gen_seq)
# plot the sequence
plot_stroke(gen_seq[0],
save_name=save_path + model_type + "_seq_" +
str(best_epoch) + ".png")
k = 0
elif k > patience:
print("Best model was saved at epoch: {}".format(best_epoch))
print("Early stopping at epoch {}".format(epoch))
break
else:
k += 1
def generate_handwriting(
char_seq="hello world",
real_text="",
style_path="../app/static/mobile/style.npy",
save_path="",
app_path="",
n_samples=1,
bias=10.0,
):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
data_path = os.path.join(app_path, "../data/")
model_path = os.path.join(app_path,
"../results/synthesis/best_model_synthesis.pt")
# seed = 194
# print("seed:",seed)
# torch.manual_seed(seed)
# np.random.seed(seed)
# print(np.random.get_state())
train_dataset = HandwritingDataset(data_path, split="train", text_req=True)
prime = True
is_map = False
style = np.load(style_path, allow_pickle=True,
encoding="bytes").astype(np.float32)
# plot the sequence
# plot_stroke(style, os.path.join(save_path, "original.png"))
print("Priming text: ", real_text)
mean, std, style = data_normalization(style)
style = torch.from_numpy(style).unsqueeze(0).to(device)
# style = valid_offset_normalization(Global.train_mean, Global.train_std, style[None,:,:])
# style = torch.from_numpy(style).to(device)
print("Priming sequence size: ", style.shape)
ytext = real_text + " " + char_seq + " "
for i in range(n_samples):
gen_seq, phi = generate_conditional_sequence(
model_path,
char_seq,
device,
train_dataset.char_to_id,
train_dataset.idx_to_char,
bias,
prime,
style,
real_text,
is_map,
)
if is_map:
plt.imshow(phi, cmap="viridis", aspect="auto")
plt.colorbar()
plt.xlabel("time steps")
plt.yticks(np.arange(phi.shape[0]),
list(ytext),
rotation="horizontal")
plt.margins(0.2)
plt.subplots_adjust(bottom=0.15)
plt.show()
# denormalize the generated offsets using train set mean and std
print("data denormalization...")
end = style.shape[1]
# gen_seq[:,:end] = data_denormalization(mean, std, gen_seq[:, :end])
# gen_seq[:,end:] = data_denormalization(Global.train_mean, Global.train_std, gen_seq[:,end:])
gen_seq = data_denormalization(Global.train_mean, Global.train_std,
gen_seq)
# plot the sequence
print(gen_seq.shape)
# plot_stroke(gen_seq[0, :end])
plot_stroke(gen_seq[0],
os.path.join(save_path, "gen_stroke_" + str(i) + ".png"))
print(save_path)
def train(model, train_loader, valid_loader, batch_size, n_epochs, lr,
patience, step_size, device, model_type, save_path):
model_path = save_path + "model_" + model_type + ".pt"
model = model.to(device)
if os.path.isfile(model_path):
model.load_state_dict(torch.load(model_path))
print(f"[ACTION] Loaded model weights from '{model_path}'")
else:
print("[INFO] No saved weights found, training from scratch.")
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = StepLR(optimizer, step_size=step_size, gamma=0.1)
train_losses = []
valid_losses = []
best_loss = math.inf
best_epoch = 0
k = 0
for epoch in range(n_epochs):
start_time = time.time()
print(f"[Epoch {epoch + 1}/{n_epochs}]")
print("[INFO] Training Model.....")
train_loss = train_epoch(model, optimizer, epoch, train_loader, device,
model_type)
print("[INFO] Validating Model....")
valid_loss = validation(model, valid_loader, device, epoch, model_type)
train_losses.append(train_loss)
valid_losses.append(valid_loss)
print(f"[RESULT] Epoch {epoch + 1}/{n_epochs}"
f"\tTrain loss: {train_loss:.3f}\tVal loss: {valid_loss:.3f}")
if step_size != -1:
scheduler.step()
if valid_loss < best_loss:
best_loss = valid_loss
best_epoch = epoch + 1
print("[SAVE] Saving weights at epoch: {}".format(epoch + 1))
torch.save(model.state_dict(), model_path)
if model_type == "prediction":
gen_seq = generate_unconditional_seq(model_path,
700,
device,
bias=10.0,
style=None,
prime=False)
else:
gen_seq = generate_conditional_sequence(
model_path,
"Hello world!",
device,
train_loader.dataset.char_to_id,
train_loader.dataset.idx_to_char,
bias=10.0,
prime=False,
prime_seq=None,
real_text=None)
# denormalize the generated offsets using train set mean and std
gen_seq = data_denormalization(Global.train_mean, Global.train_std,
gen_seq)
# plot the sequence
plot_stroke(
gen_seq[0],
save_name=save_path + model_type + "_seq_" + str(best_epoch) +
".png",
)
k = 0
elif k > patience:
print("Best model was saved at epoch: {}".format(best_epoch))
print("Early stopping at epoch {}".format(epoch))
break
else:
k += 1
total_time_taken = time.time() - start_time
print('Time taken per epoch: {:.2f}s\n'.format(total_time_taken))