def train_lstm(lstm, dataset_name, max_iter=1000, load_path=None):
dataset = LSTMDataset(name=dataset_name)
dataloader = DataLoader(dataset, batch_size=1, collate_fn=collate_fn)
if load_path is not None:
load_model(load_path, lstm)
n_iter = 0
while n_iter < max_iter:
train_loss = []
for idx, (encoded, actions) in enumerate(dataloader):
x = {'encoded': encoded, 'actions': actions}
loss = train_batch(lstm, x)
train_loss.append(loss)
if n_iter % 5 == 0:
print("[TRAIN] current iteration: {}, loss: {}".format(
n_iter, loss))
if (n_iter + 1) % 500 == 0:
dir_path = './saved_models/'
save_model(dir_path, lstm, 'lstm', str(n_iter),
str(int(time.time())), {})
n_iter += 1
print("[TRAIN] Average backward pass loss : {}".format(
np.mean(train_loss)))
def _save_model_if_eligible(self, epoch, total_score):
if total_score >= 75 and epoch % self.train_configs.save_step == 0:
save_name = "model_epoch{0}_score_{1}.pth".format(
epoch, int(total_score))
self.logger.info("Saving model as %s", save_name)
model_path = os.path.join(self.save_folder, save_name)
train_utils.save_model(model_path, self.model, self.optimizer,
epoch, total_score)
def train(self, train_loader, eval_loader):
for epoch in range(self.train_configs.start_epoch,
self.train_configs.number_of_epochs):
self.logger.info(
"Training For Epoch: %d\tLearning rate = %.4f",
epoch,
self.scheduler.get_last_lr()[0],
)
epoch_start_time = time.time()
train_size = len(train_loader.dataset)
total_loss, total_score = self._train_epoch(train_loader)
# Update learning rate. Skip updating in the last iteration.
if epoch != self.train_configs.number_of_epochs - 1:
self.scheduler.step()
total_loss /= train_size
total_score = 100 * total_score / train_size
eval_score = 0
self.logger.info(
"epoch %d,\t"
"train_size: %d,\t"
"time: %.2f,\t"
"train_loss: %.2f\t"
"SCORE: %.4f\n\n",
epoch,
train_size,
time.time() - epoch_start_time,
total_loss,
total_score,
)
if epoch == self.train_configs.number_of_epochs - 1:
self.logger.info("Saving model as %s", "final.pth")
model_path = os.path.join(self.save_folder, "final")
train_utils.save_model(model_path, self.model, self.optimizer,
epoch, total_score)
self._save_model_if_eligible(epoch, total_score)
if (eval_loader
and total_score > self.train_configs.save_score_threshold):
self.model.train(False)
self.logger.info("Threshold reached. Evaluating..")
eval_score, _ = evaluate(self.model, eval_loader)
self.model.train(True)
self.logger.info("EVAL SCORE : %.4f\n\n", eval_score * 100)
num_workers=cfg.num_thread)
R_optimizer = get_optimizer(refiner, mode='train', model_type='R')
for epoch in range(cfg.start_epoch, cfg.num_epoch):
for i, (imgs, label_maps, bg_imgs) in enumerate(Dance_dataloader):
real_imgs, label_maps, bg_imgs = imgs.cuda(), label_maps.cuda(
), bg_imgs.cuda()
input = torch.cat([label_maps, bg_imgs], dim=1)
_, fake_imgs = model(input)
refined_imgs = refiner(fake_imgs)
loss_L1 = L1_loss(real_imgs, fake_imgs)
loss_MSSSIM = MSSSIM_loss(real_imgs, fake_imgs)
R_loss = loss_MSSSIM + loss_L1
R_optimizer.zero_grad()
R_loss.backward()
R_optimizer.step()
global_logger.info(
f"{epoch}/{cfg.num_epoch} epoch, iter: {i}/{int(Dance_dataset.__len__()/cfg.batch_size)}, R_loss: {R_loss.detach().item()}"
)
R_state = {
'epoch': epoch,
'network': refiner.state_dict(),
'optimizer': R_optimizer.state_dict()
}
save_model(R_state, epoch, model_type='R')
G_loss = 4 * loss_G_GAN + 2.0 * loss_G_feat + loss_G_perceptual + 0.1 * (
loss_G_L1 +
loss_G_SmoothL1) + loss_G_intermediate + 3.0 * loss_G_MSSSIM
D_loss = 0.5 * (loss_D_fake + loss_D_real)
G_optimizer.zero_grad()
G_loss.backward()
G_optimizer.step()
D_optimizer.zero_grad()
D_loss.backward()
D_optimizer.step()
global_logger.info(
f"{epoch}/{cfg.num_epoch} epoch, iter: {i}/{int(Dance_dataset.__len__()/cfg.batch_size)}, G_loss: {G_loss.detach().item()} , D_loss: {D_loss.detach().item()}"
)
G_state = {
'epoch': epoch,
'network': model.state_dict(),
'optimizer': G_optimizer.state_dict()
}
D_state = {
'epoch': epoch,
'network': discriminator.state_dict(),
'optimizer': D_optimizer.state_dict()
}
save_model(G_state, epoch, model_type='G')
save_model(D_state, epoch, model_type='D')
def main():
# params
path_data = './data/'
path_data_samples = path_data + 'SamplesReduced/'
path_train = './train/'
path_error = './error/'
rate_testset = 0.2 # rate for the testset. All_files * rate
batch_size_train = 16
batch_size_val = 8
epochs = 2
learning_rate = 1e-4
learning_rate_decay = 1e-6
threshold_error = 20 # during testing, if [threshold_error] frames are wrong predictd, then record as error
data_shape_in_batch = [60, 75] # input shape
if not os.path.exists(path_train):
os.mkdir(path_train)
# get all the label names
list_label_names, num_output_labels = data_generator.get_labels_list(
os.path.join(path_data, 'label.txt'))
# get file list and corresponding labels
file_list_info = data_generator.get_file_list(path_data_samples,
shuffle=True)
# split the data to train set and val set
number_train = len(file_list_info) - int(
len(file_list_info) * rate_testset)
TrainInfo, ValInfo = file_list_info[0:number_train -
1], file_list_info[number_train:]
# get the data generator
gen_data_train = data_generator.data_generator_lstm(
TrainInfo,
shape=data_shape_in_batch,
batch_size=batch_size_train,
one_hot=True,
num_classes=num_output_labels,
return_sample_weights=False)
gen_data_val = data_generator.data_generator_lstm(
ValInfo,
shape=data_shape_in_batch,
batch_size=batch_size_val,
one_hot=True,
num_classes=num_output_labels,
return_sample_weights=False)
# generate keras model
model = build_model_bilstm_sequence(
num_output=num_output_labels,
input_shape=data_shape_in_batch,
learning_rate=learning_rate,
learning_rate_decay=learning_rate_decay)
model.summary()
# load model
# custom_objects = {'CRF': CRF,
# 'crf_loss': crf_loss,
# 'crf_viterbi_accuracy': crf_viterbi_accuracy}
# model = load_model(os.path.join(path_train, 'posture_recognition.h5'), custom_objects=custom_objects)
# Train
model_history = model.fit_generator(
gen_data_train,
validation_data=gen_data_val,
epochs=epochs,
steps_per_epoch=len(TrainInfo) // batch_size_train + 1,
validation_steps=len(ValInfo) // batch_size_val + 1,
verbose=1)
# save model file
train_utils.save_model(model,
save_dir=path_train,
model_base_name='posture_recognition')
# save curve
train_utils.save_model_history(model_history, 'loss', save_path=path_train)
train_utils.save_model_history(model_history,
'crf_viterbi_accuracy',
save_path=path_train,
show_test_value=True)
# save confusion matrix
gen_data_train = data_generator.data_generator_lstm(
TrainInfo,
shape=data_shape_in_batch,
batch_size=batch_size_train,
one_hot=False,
num_classes=num_output_labels,
is_training=False,
return_sample_weights=False)
gen_data_val = data_generator.data_generator_lstm(
ValInfo,
shape=data_shape_in_batch,
batch_size=batch_size_val,
one_hot=False,
num_classes=num_output_labels,
is_training=False,
return_sample_weights=False)
train_utils.predict_from_generator_save_error(
model,
gen_data_train,
step=len(TrainInfo) // batch_size_train + 1,
threshold=threshold_error,
save_path_error=path_error,
save_cm=True,
save_path_cm=path_train,
list_label_names=list_label_names,
title='train')
train_utils.predict_from_generator_save_error(
model,
gen_data_val,
step=len(ValInfo) // batch_size_val + 1,
threshold=threshold_error,
save_path_error=path_error,
save_cm=True,
save_path_cm=path_train,
list_label_names=list_label_names,
title='val')
def main():
# ----- get config -----
config = parse_args_get_config()
# ----- logger -----
logger = SummaryWriter(log_dir=config.log_path + "/" + config.run_name)
for c, c_value in config.items():
logger.add_text(c, str(c_value))
# ----- set seed -----
set_seed(config)
# ----- get device -----
device, n_gpu = get_device_info(config)
config.device = device
# ----- load data object -----
print("Loading data")
data = get_data(config)
config['n_updates_total'] = len(data.train_iter) * config.epochs
# ----- create or load model -----
print("Loading model")
model = get_model(config, data.vocab)
if config.restore_model:
print("restoring")
restore_model(model, config.restore_path, device)
model.to(device)
pytorch_total_params = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print("Total params: ", pytorch_total_params)
print(model)
# ----- create criterion -----
print("Loading criterion")
criterion = get_criterion(config)
# ----- create optimizer -----
print("Loading optimizer")
optimizer = get_optimizer(config, model)
# ----- get loss compute function -----
print("Loading loss compute function")
loss_compute = get_loss_compute(config, criterion, optimizer)
loss_compute_dev = get_loss_compute(config, criterion, None)
# ----- train mode -----
if config.mode == 'train':
print("Training")
best_dev_acc = -1
for i in range(config.epochs):
# ----- train -----
model.train()
run_epoch(logger,
config,
i,
data,
data.train_iter,
model,
loss_compute,
device,
mode='train',
save_misclassified=False)
# ----- dev -----
model.eval()
with torch.no_grad():
dev_acc, dev_loss = run_epoch(logger,
config,
i,
data,
data.dev_iter,
model,
loss_compute_dev,
device,
mode='dev',
save_misclassified=False)
if dev_acc > best_dev_acc:
best_dev_acc = dev_acc
if config.save_model:
save_model(model, config, dev_acc, i)
# ----- test -----
with torch.no_grad():
run_epoch(logger,
config,
i,
data,
data.test_iter,
model,
loss_compute_dev,
device,
mode='test',
save_misclassified=True)
# ----- test mode -----
elif config.mode == 'test':
print("Testing")
model.eval()
with torch.no_grad():
test_acc, test_loss = run_epoch(logger,
config,
0,
data,
data.test_iter,
model,
loss_compute_dev,
device,
mode='test',
save_misclassified=True)
print("Test loss: ", test_loss)
print("Test acc: ", test_acc)
# ----- interactive mode -----
elif config.mode == 'interactive':
print("Interactive")
model.eval()
while True:
premise = input("premise > ")
hypothesis = input("hypothesis > ")
if not premise or not hypothesis:
print("Please enter premise and hypothesis")
continue
premise_var = text_to_var(premise, data.TEXT, config)
hypothesis_var = text_to_var(hypothesis, data.TEXT, config)
mbatch = dotdict({
'premise': premise_var,
'hypothesis': hypothesis_var,
})
sent_p = premise.split()
sent_h = hypothesis.split()
scores = model(mbatch)
model.draw(sent_p, sent_h)
logger.close()