def setUp(self): self.indir, self.workdir, self.outdir = setup_dirs(fpath=__file__) torch.manual_seed(0) self.model = LSTMModel(sizes=[4, 128, 2], bidirectional=True, softmax_output=True, cuda=False) X_train = (torch.rand(100, 32, 4) > 0.5).float() # (n_samples, seq_len, input_size) Y_train = (torch.rand(100, 32) > 0.5).long() # (n_samples, seq_len) X_test = (torch.rand(50, 32, 4) > 0.5).float() Y_test = (torch.rand(50, 32) > 0.5).long() criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(self.model.parameters()) self.trainer = Trainer(model=self.model, X_train=X_train, Y_train=Y_train, X_test=X_test, Y_test=Y_test, criterion=criterion, optimizer=optimizer, accuracy=softmax_accuracy, mini_batch_size=16, log_dir=f'{self.workdir}/tensorboard/test')
def main(): gbkdir = '../data' gbks = get_files(source=gbkdir, isfullpath=True) X, Y = load_genbank(gbks=gbks, cuda=False) model = LSTMModel(sizes=[4, 128, 64, 1], batch_first=True, bidirectional=True, cuda=CUDA) for stage, seq_len, mini_batch_size, learning_rate in [ (1, 128, 256, 1e-3), (2, 1024, 32, 1e-4), (3, 8192, 4, 1e-5), ]: model = train_model(stage=stage, X=X, Y=Y, seq_len=seq_len, mini_batch_size=mini_batch_size, learning_rate=learning_rate, model=model) torch.save(model, f'./models/{EXPERIMENT_NAME}_stage_{stage}.model') torch.save(model, f'./models/{EXPERIMENT_NAME}.model')
def train_model(stage: int, X: torch.Tensor, Y: torch.Tensor, seq_len: int, mini_batch_size: int, learning_rate: float, model: LSTMModel) -> LSTMModel: X, Y = divide_sequence(X=X, Y=Y, seq_len=seq_len) X, Y = shuffle(X=X, Y=Y) X_train, X_test = split(X, training_fraction=TRAINING_FRACTION) Y_train, Y_test = split(Y, training_fraction=TRAINING_FRACTION) criterion = nn.BCEWithLogitsLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) writer = SummaryWriter( log_dir= f'tensorboard/{EXPERIMENT_NAME}/stage_{stage}_seq_len_{seq_len}_lr_{learning_rate}' ) trainer = Trainer(model=model, X_train=X_train, Y_train=Y_train, X_test=X_test, Y_test=Y_test, criterion=criterion, optimizer=optimizer, mini_batch_size=mini_batch_size, writer=writer) model = trainer.train(max_epochs=MAX_EPOCHS, overtrain_epochs=OVERTRAIN_EPOCHS) writer.close() return model
def main(): for i, architecture in enumerate(ARCHITECTURES): model = LSTMModel( sizes=architecture, bidirectional=BIDIRECTIONAL, softmax_output=SOFTMAX, cuda=CUDA) for stage, mini_batch_size, seq_len, learning_rate in TRAINING_STAGES: model = train_model( gbkdir=GBKDIR, stage=stage, mini_batch_size=mini_batch_size, seq_len=seq_len, learning_rate=learning_rate, model=model, model_name=f'model_{i+1}') os.makedirs('./models', exist_ok=True) torch.save(model, f'./models/{EXPERIMENT_NAME}_model_{i+1}_stage_{stage}.model') torch.save(model, f'./models/{EXPERIMENT_NAME}_model_{i+1}.model')
def setUp(self): self.indir, self.workdir, self.outdir = setup_dirs(fpath=__file__) self.model = LSTMModel(sizes=[4, 128, 64, 32, 2], bidirectional=True, softmax_output=True, cuda=False)
def test___init__(self): expected = f'''\ LSTMModel( (lstm_layers): ModuleList( (0): LSTM(4, 128, batch_first=True, bidirectional=True) (1): LSTM(256, 64, batch_first=True, bidirectional=True) (2): LSTM(128, 32, batch_first=True, bidirectional=True) ) (fc): Linear(in_features=64, out_features=2, bias=True) )''' self.assertEqual(expected, str(self.model))
def load_model(file: str, cuda: bool) -> LSTMModel: model = torch.load(file, map_location=torch.device('cpu')) # Because source code was changed, need to unpack the parameters and # load them into a newly instantiated LSTMModel state_dict = {} for key, val in model.state_dict().items(): k = key.replace('_1', '_layers.0').replace('_2', '_layers.1') state_dict[k] = val new_model = LSTMModel( sizes=[4, 128, 64, 1], bidirectional=True, softmax_output=False, cuda=cuda) new_model.load_state_dict(state_dict) return new_model
def train_model( gbkdir: str, stage: int, mini_batch_size: int, seq_len: int, learning_rate: float, model: LSTMModel, model_name: str) -> LSTMModel: gbks = get_files(source=gbkdir, isfullpath=True) X, Y = load_genbank(gbks=gbks, label_length=LABEL_LENGTH) if not SOFTMAX: # Reshape for binary classification Y = Y.view(-1, 1).float() # 1D long -> 2D float X = divide_sequence(X, seq_len=seq_len, pad=True) Y = divide_sequence(Y, seq_len=seq_len, pad=True) X, Y = shuffle(X, Y) X_train, X_test = split(X, training_fraction=TRAINING_FRACTION, dim=0) Y_train, Y_test = split(Y, training_fraction=TRAINING_FRACTION, dim=0) weight = get_class_weight(Y) if CUDA: weight = weight.cuda() criterion = nn.CrossEntropyLoss(weight=weight) if SOFTMAX else nn.BCEWithLogitsLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) log_dir = f'tensorboard/{EXPERIMENT_NAME}/{model_name}_stage_{stage}' accuracy = softmax_accuracy if SOFTMAX else binary_accuracy trainer = Trainer( model=model, X_train=X_train, Y_train=Y_train, X_test=X_test, Y_test=Y_test, criterion=criterion, optimizer=optimizer, accuracy=accuracy, mini_batch_size=mini_batch_size, log_dir=log_dir) model = trainer.train( max_epochs=MAX_EPOCHS, overtrain_epochs=OVERTRAIN_EPOCHS) return model