base_model.py 648 B

12345678910111213141516171819202122232425
  1. import torch.nn as nn
  2. import numpy as np
  3. from abc import abstractmethod
  4. class BaseModel(nn.Module):
  5. """
  6. Base class for all models
  7. """
  8. @abstractmethod
  9. def forward(self, *inputs):
  10. """
  11. Forward pass logic
  12. :return: Model output
  13. """
  14. raise NotImplementedError
  15. def __str__(self):
  16. """
  17. Model prints with number of trainable parameters
  18. """
  19. model_parameters = filter(lambda p: p.requires_grad, self.parameters())
  20. params = sum([np.prod(p.size()) for p in model_parameters])
  21. return super().__str__() + '\nTrainable parameters: {}'.format(params)