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