File size: 3,153 Bytes
bc45d7d 467ec0d bc45d7d c7bd2b8 bc45d7d 263280b bc45d7d 263280b 693ac2f 9589c44 bc45d7d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | import warnings
import argparse
import train_nn
import train_xgb
from cnn_runner import save_cnn_features
parser = argparse.ArgumentParser(
description="INCLUDE trainer for xgboost, lstm and transformer"
)
parser.add_argument("--seed", default=0, type=int, help="seed value")
parser.add_argument(
"--dataset",
default="isl_split_dataset",
type=str,
help="dataset prefix used by label map and split folders",
)
parser.add_argument(
"--use_augs",
action="store_true",
help="use augmented data",
)
parser.add_argument(
"--use_cnn",
action="store_true",
help="use mobilenet to convert keypoints to videos and generate embeddings from CNN",
)
parser.add_argument(
"--model",
default="lstm",
type=str,
help="options: lstm, transformer, xgboost",
)
parser.add_argument(
"--data_dir",
default="",
type=str,
required=True,
help="location to train, val and test json files",
)
parser.add_argument(
"--save_path",
default="./",
type=str,
help="location to save trained model",
)
parser.add_argument(
"--epochs", default=150, type=int, help="number of epochs to train the model"
)
parser.add_argument("--batch_size", default=128, type=int, help="batch size of data")
parser.add_argument(
"--learning_rate",
default=1e-4,
type=float,
help="learning rate for training neural net",
)
parser.add_argument(
"--transformer_size", default="small", type=str, help="options: small, large"
)
parser.add_argument(
"--max_frame_len",
default=169,
type=int,
help="sequence length for train/eval keypoint padding",
)
parser.add_argument(
"--use_pretrained",
default=None,
help="use pretrained model. options: evaluate, resume_training",
)
parser.add_argument(
"--eval_split",
default="test",
choices=["train", "val", "test"],
help="split to evaluate when running evaluate",
)
parser.add_argument(
"--early_stop_patience",
default=15,
type=int,
help="epochs with no improvement before early stopping",
)
parser.add_argument(
"--early_stop_metric",
default="val_acc",
choices=["val_loss", "val_acc", "loss_gap"], # add loss_gap
help="metric used for checkpointing and early stopping during training",
)
args = parser.parse_args()
if __name__ == "__main__":
if args.model == "xgboost":
if args.use_pretrained:
raise Exception("Pre-trained models are not available for XGBoost")
if args.use_cnn:
warnings.warn(
"use_cnn flag set to true for xgboost model. xgboost will not use cnn features"
)
train_xgb.fit(args)
train_xgb.evaluate(args)
else:
if args.use_cnn:
save_cnn_features(args)
if args.use_augs:
warnings.warn("cannot perform augmentation on cnn features")
if args.use_pretrained == "evaluate":
train_nn.evaluate(args)
print("### Evaluated from pretrained model ###")
else:
print("### Starting to train. ###")
train_nn.fit(args)
train_nn.evaluate(args)
|