def __init__(self, T, loss='mine', alpha=0.01, method=None):
super().__init__()
self.running_mean = 0
self.loss = loss
self.alpha = alpha
self.method = method
if method == 'concat':
if isinstance(T, nn.Sequential):
self.T = CustomSequential(ConcatLayer(), *T)
else:
self.T = CustomSequential(ConcatLayer(), T)
else:
self.T = T
def __init__(self, x_dim, z_dim):
super().__init__()
self.layers = CustomSequential(ConcatLayer(),
nn.Linear(x_dim + z_dim,
400), nn.ReLU(),
nn.Linear(400, 400), nn.ReLU(),
nn.Linear(400, 400), nn.ReLU(),
nn.Linear(400, 1))
def __init__(self, x_dim, z_dim, loss='mine', **kwargs):
super().__init__()
self.x_dim = x_dim
self.T = CustomSequential(ConcatLayer(), nn.Linear(x_dim + z_dim, 100),
nn.ReLU(), nn.Linear(100, 100), nn.ReLU(),
nn.Linear(100, 1))
self.energy_loss = Mine(self.T, loss=loss, alpha=kwargs['alpha'])
self.kwargs = kwargs
self.train_loader = kwargs.get('train_loader')
self.test_loader = kwargs.get('test_loader')