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
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 test_divide_sequence(self):
x = torch.randn(100, 4)
for pad, expected in [(True, (34, 3, 4)), (False, (33, 3, 4))]:
size = divide_sequence(x, seq_len=3, pad=pad).size()
self.assertTupleEqual(expected, size)
def __forward_pass(self, X: torch.Tensor) -> torch.Tensor:
"""
1) Break the input X into segments, or windows
2) Forward pass to get Y, i.e. predicted probability
3) Flatten windows of Y back to full input length
Flat-in, flat-out
Args:
X: 2D tensor, size (full_len, 4), dtype=torch.float
Returns:
Y: 1D tensor, size (full_len, ), dtype=torch.float
"""
full_len = X.shape[0]
X = divide_sequence(X, seq_len=self.window, pad=True) # 2D -> 3D
n_samples, window, _ = X.size()
Y = torch.zeros(n_samples, window) # 2D
i = 0
mb = self.mini_batch_size
while i < n_samples:
start = i
end = min(i + mb,
n_samples) # the last mini_batch < mini_batch_size
x = X[start:end, :, :]
if self.model.is_cuda:
x = x.cuda()
with torch.no_grad():
y_hat = self.model(x)
y = self.output_to_label(
y_hat) # 3D -> 2D tensor (batch_size, window)
Y[start:end, :] = y[:, :]
i += mb
Y = Y.view(-1) # 2D (n_samples, seq_len) -> 1D (full_len)
Y = Y[0:full_len] # remove padded sequence
return Y.float()
def validate_model():
"""
Training with seq_len = 8192 kept failing
I just want to see how the model performs on such a long sequence
"""
gbkdir = '../data'
seq_len = 8192
mini_batch_size = 32
gbks = get_files(source=gbkdir, isfullpath=True)
X, Y = load_genbank(gbks=gbks, cuda=False)
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)
model = torch.load(f'./models/{EXPERIMENT_NAME}_stage_2.model')
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters())
writer = SummaryWriter(
log_dir=f'tensorboard/{EXPERIMENT_NAME}/validate_seq_len_{seq_len}')
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)
trainer.validate()
writer.close()