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()