| 87 | |
| 88 | |
| 89 | class ComputeCovA: |
| 90 | |
| 91 | @classmethod |
| 92 | def compute_cov_a(cls, a, layer): |
| 93 | return cls.__call__(a, layer) |
| 94 | |
| 95 | @classmethod |
| 96 | def __call__(cls, a, layer): |
| 97 | if isinstance(layer, nn.Linear): |
| 98 | cov_a = cls.linear(a, layer) |
| 99 | elif isinstance(layer, nn.Conv2d): |
| 100 | cov_a = cls.conv2d(a, layer) |
| 101 | else: |
| 102 | # FIXME(CW): for extension to other layers. |
| 103 | # raise NotImplementedError |
| 104 | cov_a = None |
| 105 | |
| 106 | return cov_a |
| 107 | |
| 108 | @staticmethod |
| 109 | def conv2d(a, layer): |
| 110 | batch_size = a.size(0) |
| 111 | a = _extract_patches(a, layer.kernel_size, layer.stride, layer.padding) |
| 112 | spatial_size = a.size(1) * a.size(2) |
| 113 | a = a.view(-1, a.size(-1)) |
| 114 | if layer.bias is not None: |
| 115 | a = torch.cat([a, a.new(a.size(0), 1).fill_(1)], 1) |
| 116 | a = a/spatial_size |
| 117 | # FIXME(CW): do we need to divide the output feature map's size? |
| 118 | return a.t() @ (a / batch_size) |
| 119 | |
| 120 | @staticmethod |
| 121 | def linear(a, layer): |
| 122 | # a: batch_size * in_dim |
| 123 | batch_size = a.size(0) |
| 124 | if layer.bias is not None: |
| 125 | a = torch.cat([a, a.new(a.size(0), 1).fill_(1)], 1) |
| 126 | return a.t() @ (a / batch_size) |
| 127 | |
| 128 | |
| 129 | class ComputeCovG: |