Skip to content

Commit

Permalink
Merge pull request #864 from marrlab/causalIRL
Browse files Browse the repository at this point in the history
first version causal irl
  • Loading branch information
smilesun authored Aug 2, 2024
2 parents 4301cd4 + 2c1cdc9 commit 88fd530
Show file tree
Hide file tree
Showing 3 changed files with 94 additions and 1 deletion.
77 changes: 77 additions & 0 deletions domainlab/algos/trainers/train_causIRL.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,77 @@
"""
Alex, Xudong
"""
import numpy as np
import torch
from domainlab.algos.trainers.train_basic import TrainerBasic


class TrainerCausalIRL(TrainerBasic):
"""
causal matching
"""
def my_cdist(self, x1, x2):
"""
distance for Gaussian
"""
# along the last dimension
x1_norm = x1.pow(2).sum(dim=-1, keepdim=True)
x2_norm = x2.pow(2).sum(dim=-1, keepdim=True)
# x_2_norm is [batchsize, 1]
# matrix multiplication (2nd, 3rd) and addition to first argument
# X1[batchsize, dimfeat] * X2[dimfeat, batchsize)
# alpha: Scaling factor for the matrix product (default: 1)
# x2_norm.transpose(-2, -1) is row vector
# x_1_norm is column vector
res = torch.addmm(x2_norm.transpose(-2, -1),
x1,
x2.transpose(-2, -1), alpha=-2).add_(x1_norm)
return res.clamp_min_(1e-30)

def gaussian_kernel(self, x, y):
"""
kernel for MMD
"""
gamma=[0.001, 0.01, 0.1, 1, 10, 100, 1000]
dist = self.my_cdist(x, y)
tensor = torch.zeros_like(dist)
for g in gamma:
tensor.add_(torch.exp(dist.mul(-g)))
return tensor

def mmd(self, x, y):
"""
maximum mean discrepancy
"""
kxx = self.gaussian_kernel(x, x).mean()
kyy = self.gaussian_kernel(y, y).mean()
kxy = self.gaussian_kernel(x, y).mean()
return kxx + kyy - 2 * kxy

def tr_batch(self, tensor_x, tensor_y, tensor_d, others, ind_batch, epoch):
"""
optimize neural network one step upon a mini-batch of data
"""
self.before_batch(epoch, ind_batch)
tensor_x, tensor_y, tensor_d = (
tensor_x.to(self.device),
tensor_y.to(self.device),
tensor_d.to(self.device),
)
self.optimizer.zero_grad()

features = self.get_model().extract_semantic_feat(tensor_x)

pos_batch_break = np.random.randint(0, tensor_x.shape[0])
first = features[:pos_batch_break]
second = features[pos_batch_break:]
if len(first) > 1 and len(second) > 1:
penalty = torch.nan_to_num(self.mmd(first, second))
else:
penalty = torch.tensor(0)
loss = self.cal_loss(tensor_x, tensor_y, tensor_d, others)
loss = loss + penalty
loss.backward()
self.optimizer.step()
self.after_batch(epoch, ind_batch)
self.counter_batch += 1
5 changes: 4 additions & 1 deletion domainlab/algos/trainers/zoo_trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,13 @@
from domainlab.algos.trainers.train_basic import TrainerBasic
from domainlab.algos.trainers.train_ema import TrainerMA
from domainlab.algos.trainers.train_dial import TrainerDIAL
from domainlab.algos.trainers.train_hyper_scheduler import TrainerHyperScheduler
from domainlab.algos.trainers.train_hyper_scheduler \
import TrainerHyperScheduler
from domainlab.algos.trainers.train_matchdg import TrainerMatchDG
from domainlab.algos.trainers.train_mldg import TrainerMLDG
from domainlab.algos.trainers.train_fishr import TrainerFishr
from domainlab.algos.trainers.train_irm import TrainerIRM
from domainlab.algos.trainers.train_causIRL import TrainerCausalIRL


class TrainerChainNodeGetter(object):
Expand Down Expand Up @@ -54,6 +56,7 @@ def __call__(self, lst_candidates=None, default=None, lst_excludes=None):
chain = TrainerFishr(chain)
chain = TrainerIRM(chain)
chain = TrainerHyperScheduler(chain)
chain = TrainerCausalIRL(chain)
node = chain.handle(self.request)
head = node
while self._list_str_trainer:
Expand Down
13 changes: 13 additions & 0 deletions tests/test_causal_irl.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
"""
end-end test
"""
from tests.utils_test import utils_test_algo


def test_causal_irl():
"""
causal irl
"""
args = "--te_d 0 --tr_d 3 7 --bs=32 --debug --task=mnistcolor10 \
--model=erm --nname=conv_bn_pool_2 --trainer=causalirl"
utils_test_algo(args)

0 comments on commit 88fd530

Please sign in to comment.