<jupyter_start><jupyter_text>Lightning model<jupyter_code>class Net(pl.LightningModule):
    """Pytorch Lightning Module for defining and training the model
    
    This is a Pytorch Lightning clas which helps us to rewrite the pytorch code without the boilerplate 
    codes
    """
    
    def __init__(self, model, val_image_count=6):
        """Init the model
        
        Args:
            model: Model to train
            val_image_count: No of images to log during validation
        """
        
        super().__init__()
        self.model = model
        
        self.val_image_count = val_image_count
        
        # Saving the best model
        self.lowest_val_acc = 0
        self.best_epoch = None
        
    def forward(self, x):
        """Forward method for the net
        Args:
            x: Any tensor to input in the model
        """
        
        return self.model(x)
    
    def training_step(self, batch, batch_idx):
        """Training step
        
        Args:
            batch: Batches of input data from train_dataloader
            batch_idx: Index of the Batch or the step count
        
        Returns:
            A dict containing loss and logs for tensorboard
        """
        
        x, y = batch
        yt = self(x)
        loss = F.cross_entropy(yt, y)
        tensorboard_logs = {'train_loss': loss}
        return {'loss': loss, 'log': tensorboard_logs}
    
    def validation_step(self, batch, batch_idx):
        """Validation step
        
        Args:
            batch: Batches of input data from val_dataloader
            batch_idx: Index of the Batch or the step count
        
        Returns:
            A dict containing validation loss and validation accuracy
        """
        
        x, y = batch

        yt = self(x)
        loss = F.cross_entropy(yt, y)
        
        preds = torch.argmax(yt, dim=1)
        val_acc = torch.sum(y == preds).item() / (len(y) * 1.0)
#         self.logger.experiment.add_graph(self.model, x)
        if batch_idx == 0:
            probs = [F.softmax(el, dim=0)[i].item() for i, el in zip(preds[:self.val_image_count], 
                                                                     yt[:self.val_image_count])]
            self.logger.experiment.add_figure('predictions vs. actuals',
                                              self.plot_classes_preds(preds[:self.val_image_count], 
                                                                      probs, x[:self.val_image_count], 
                                                                      y[:self.val_image_count]),
                                              global_step=self.current_epoch)
        
        output = {
            'val_loss': loss,
            'val_acc': torch.tensor(val_acc),
        }

        return output
    
    def validation_epoch_end(self, outputs):
        """After all the Validation steps are completed
        
        Args:
            outputs: All the outputs of validation steps
        Returns:
            results: A dict containing progress bar update and tensorboard logs
        """
        
        val_acc_mean = 0
        for output in outputs:
            val_acc_mean += output['val_acc']

        val_acc_mean /= len(outputs)
        tqdm_dict = {'val_acc': val_acc_mean.item()}
        
        # Saving_best_epoch
        if val_acc_mean.item() > self.lowest_val_acc:
            self.lowest_val_acc = val_acc_mean.item()
            self.best_epoch = self.current_epoch
            torch.save(self.model.state_dict(), './../Models/best_model.pth')

        results = {
            'progress_bar': tqdm_dict,
            'log': {'val_acc': val_acc_mean, 'epoch': self.current_epoch, 'best_epoch': self.best_epoch}
        }
        return results
    
    def plot_classes_preds(self, preds, probs, images, labels):
        """Plot the images in a matplotlib figure
        
        Args:
            preds: Predicted label ids of images
            probs: Prediction probabilities of the class
            images: Images to be plotted
            labels: Labels corresponding to preds
        
        Returns:
            fig: The matplotlip figure containing subplots of image and labels
        """
        
        fig = plt.figure(figsize=(12, 12))
        for idx in range(images.shape[0]):
            ax = fig.add_subplot(1, images.shape[0], idx+1, xticks=[], yticks=[])
            matplotlib_imshow(images[idx], one_channel=True)
            ax.set_title("{0}, {1:.1f}%\n(label: {2})".format(
                self.classes[preds[idx]],
                probs[idx] * 100.0,
                self.classes[labels[idx]]),
                        color=("green" if preds[idx]==labels[idx].item() else "red"))
        return fig

    def configure_optimizers(self):
        """Configuring the optimizers to use for training"""
        
        # return torch.optim.Adam(self.parameters(), lr=0.001)
        return torch.optim.RMSprop(self.parameters(), lr=0.0001, weight_decay=1e-6)

    def train_dataloader(self):
        """Dataloader for training"""
        
        train = get_data(train=True)
        loader = DataLoader(train, batch_size=32, num_workers=4, shuffle=True, pin_memory=True)
        return loader
    
    def val_dataloader(self):
        """Dataloader for validation"""
        
        val = get_data()
        self.classes = val.classes
        loader = DataLoader(val, batch_size=32, num_workers=4, shuffle=False, pin_memory=True)
        return loader<jupyter_output><empty_output><jupyter_text>Training<jupyter_code># Model to train
model = Net(efficientNet())

# Checkpoint incase there is an interruption
checkpoint_callback = ModelCheckpoint(
    filepath=os.path.join(os.getcwd(), './lightning_logs/checkpoints'),
    save_top_k=True,
    verbose=True,
    monitor='val_acc',
    mode='max',
    prefix=''
)

# Trainer to train the model on all GPUs for 20 epochs
trainer = Trainer(max_epochs=20, gpus=-1, checkpoint_callback=checkpoint_callback)#, resume_from_checkpoint='./lightning_logs/Checkpoints_efficientNet/_ckpt_epoch_19.ckpt')
trainer.fit(model)   # Fit the model<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09566903684550743, 'fraction_numerical': 0.007756948933419522, 'mean_word_length': 1.8429030776297657, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 13, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '607127', 'license': 'MIT', 'n_tokens_mistral': 1817, 'n_tokens_neox': 1770, 'n_words': 435, 'path': 'Research/train.ipynb', 'repo_name': 'Bibyutatsu/TorchRT', 'size': '6188'}