Custom training (without back propagation)
What is your question?
I have a custom nn.module of Gaussian Mixture Models Classifier based on sklearn's, which is trained by E-M algorithm instead of backpropagation.
As I'm trying to use this framework to implement training, evaluations and checkpointing. I've come to the problem that I can't understand how to override the optimizer and back-propagation behavior.
As this custom classifier trains with it .fit method. I want to call this method in each train step, instead of doing a forward pass and getting the loss value and passing it to the optimizer.
Code
class GMMClassifier(nn.Module):
    def __init__(self, num_classes):
        super(GMMClassifier, self).__init__()
        n_features = 128
        self.gmm_list = []
        for _ in num_classes:
            # one gmm per singer as stated in Tsai; Fujihara; Mesaros et. al works on SID
            self.gmm_list.append(GaussianMixture(n_components=64, n_features=n_features))

    def forward_score(self, x):
        """
        :param x: MFCC of a track with shape (samples, frames, coefficients, )
        :return: The Log Likelihood for each track and frame tested on every GMM (one per singer / class) as:
            log likelihood = log p(X_t / P_i)]
            where t is time frame and
                i is the singer GMM
            with shape: (sample, frame, gmm_prediction)
        """
        n_samples = x.size(0)
        # asume that all the samples has equal frame number
        n_frames = x.size(1)
        n_features = x.size(2)
        score_tensors_per_gmm = []
        # print('info: feeding gmms...')
        for gmm in self.gmm_list:
            # predict each frame for each sampple
            # optimization: flatten (samples, frames, features) to (samples*frames, features)
            x = x.reshape(-1, n_features)
            log_prob = gmm.score(x)  # output shape is (samples*frames, )
            log_prob = torch.unsqueeze(log_prob, dim=1)  # output shape is (samples*frames, 1) for stacking
            score_tensors_per_gmm.append(
                log_prob
            )
        score_tensors_per_gmm = torch.stack(score_tensors_per_gmm,
                                            dim=1)  # output tensor shape is ((samples*frames, gmm)
        y = score_tensors_per_gmm.reshape(
            n_samples,
            n_frames,
            len(self.gmm_list)
        )  # reshape to (sample, frame, gmm)
        return y

    def forward(self, x):
        """
        :param x: MFCC of a track with shape (samples, frames, coefficients, )
        :return: The prediction for each track calculated as:
            Singer_id = arg max_i [1/T sum^T_t=1 [log p(X_t / P_i)]]
            where t is time frame and
                i is the singer GMM
            with shape: (n_samples, )
        """
        x = self.forward_score(x)  # shape is (sample, frame, gmm)
        x = x.sum(dim=1)  # shape is (sample, gmm)
        x = x.argmax(dim=1)  # shape is (sample, )
        # sum the scores over the
        return x
class L_GMMClassifier(ptl.LightningModule):
    """
    Sample model to show how to define a template
    """

    def __init__(self, hparams, num_classes, train_dataset, eval_dataset, test_dataset):
        super(L_GMMClassifier, self).__init__()
        self.hparams = hparams
        # self.loss = torch.nn.CrossEntropyLoss()
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.test_dataset = test_dataset
        # build model
        self.model = GMMClassifier(num_classes)
        # self.optimizer = torch.optim.Adam(self.model.parameters(), lr=hparams.learning_rate)

    # ---------------------
    # TRAINING
    # ---------------------
    def forward(self, x):
        """
        No special modification required for lightning, define as you normally would
        :param x:
        :return:
        """
        return self.model(x)

    def training_step(self, batch, batch_idx):
        """
        Lightning calls this inside the training loop
        :param batch:
        :return:
        """
        # forward pass
        x, y = batch['x'], batch['y']
        for gmm_idx, gmm in enumerate(self.model.gmm_list):
            gmm.fit()