Here a mask with continuous values in the range [0,1] is formed to control the amount of update for each parameter based on the agreement of gradients coming from different environments.
(self, gradients, params)
| 1145 | return {'loss': mean_loss} |
| 1146 | |
| 1147 | def mask_grads(self, gradients, params): |
| 1148 | ''' |
| 1149 | Here a mask with continuous values in the range [0,1] is formed to control the amount of update for each |
| 1150 | parameter based on the agreement of gradients coming from different environments. |
| 1151 | ''' |
| 1152 | device = gradients[0][0].device |
| 1153 | for param, grads in zip(params, gradients): |
| 1154 | grads = torch.stack(grads, dim=0) |
| 1155 | avg_grad = torch.mean(grads, dim=0) |
| 1156 | grad_signs = torch.sign(grads) |
| 1157 | gamma = torch.tensor(1.0).to(device) |
| 1158 | grads_var = grads.var(dim=0) |
| 1159 | grads_var[torch.isnan(grads_var)] = 1e-17 |
| 1160 | lam = (gamma * grads_var).pow(-1) |
| 1161 | mask = torch.tanh(self.k * lam * (torch.abs(grad_signs.mean(dim=0)) - self.tau)) |
| 1162 | mask = torch.max(mask, torch.zeros_like(mask)) |
| 1163 | mask[torch.isnan(mask)] = 1e-17 |
| 1164 | mask_t = (mask.sum() / mask.numel()) |
| 1165 | param.grad = mask * avg_grad |
| 1166 | param.grad *= (1. / (1e-10 + mask_t)) |
| 1167 | |
| 1168 | |
| 1169 | class Fishr(Algorithm): |