File size: 573 Bytes
78a947a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
"""
Contains an implementation of a Module that averages the outputs of multiple models.
This can be used to construct an ensemble of models.
Author: Ole-Christian Galbo Engstrøm
E-mail: ocge@foss.dk
"""
import torch
import torch.nn as nn
class EnsembleModel(nn.Module):
def __init__(self, *models):
super(EnsembleModel, self).__init__()
self.models = nn.ModuleList(models)
def forward(self, x):
outputs = torch.stack([model(x) for model in self.models])
return torch.mean(outputs, axis=0) # Average the outputs of all models
|