| 12345678910111213141516171819202122232425 |
- import torch.nn as nn
- import numpy as np
- from abc import abstractmethod
- class BaseModel(nn.Module):
- """
- Base class for all models
- """
- @abstractmethod
- def forward(self, *inputs):
- """
- Forward pass logic
- :return: Model output
- """
- raise NotImplementedError
- def __str__(self):
- """
- Model prints with number of trainable parameters
- """
- model_parameters = filter(lambda p: p.requires_grad, self.parameters())
- params = sum([np.prod(p.size()) for p in model_parameters])
- return super().__str__() + '\nTrainable parameters: {}'.format(params)
|