| 127 | |
| 128 | |
| 129 | def kaiming_init(m, is_linear, nonlinearity="silu"): |
| 130 | # gain = math.sqrt(2.0 / (1.0 + alpha**2)) |
| 131 | |
| 132 | # gain=np.sqrt(10.5) |
| 133 | if nonlinearity=="silu": |
| 134 | gain=np.sqrt(2.3) #works fine with silu |
| 135 | elif nonlinearity=="relu": |
| 136 | gain=np.sqrt(2) #works fine with silu |
| 137 | # gain=np.sqrt(2.15) |
| 138 | # gain=np.sqrt(0.92) #for mpsilu |
| 139 | scale=1.0 |
| 140 | |
| 141 | if is_linear: |
| 142 | gain = 1 |
| 143 | |
| 144 | # print("effective scale", gain*scale) |
| 145 | |
| 146 | # print("m is ",m) |
| 147 | # help(m) |
| 148 | |
| 149 | if isinstance(m, torch.nn.Conv2d): |
| 150 | ksize = m.kernel_size[0] * m.kernel_size[1] |
| 151 | n1 = m.in_channels |
| 152 | n2 = m.out_channels |
| 153 | |
| 154 | # std = gain * math.sqrt(2.0 / ((n1 + n2) * ksize)) |
| 155 | std = gain * math.sqrt(n1 * ksize) |
| 156 | elif isinstance(m, torch.nn.ConvTranspose2d): |
| 157 | ksize = m.kernel_size[0] * m.kernel_size[1] // 4 |
| 158 | n1 = m.in_channels |
| 159 | n2 = m.out_channels |
| 160 | |
| 161 | # std = gain * math.sqrt(2.0 / ((n1 + n2) * ksize)) |
| 162 | std = gain * math.sqrt(n1 * ksize) |
| 163 | elif isinstance(m, torch.nn.Conv1d): |
| 164 | ksize = m.kernel_size[0] |
| 165 | n1 = m.in_channels |
| 166 | n2 = m.out_channels |
| 167 | # std = gain * math.sqrt(2.0 / ((n1 + n2) * ksize)) |
| 168 | std = gain * math.sqrt(n1 * ksize) |
| 169 | elif isinstance(m, torch.nn.Linear): |
| 170 | n1 = m.in_features |
| 171 | n2 = m.out_features |
| 172 | |
| 173 | # std = gain * math.sqrt(2.0 / (n1 + n2)) |
| 174 | std = gain * math.sqrt(n1) |
| 175 | # elif isinstance(m, SMConv1d_v2): |
| 176 | # print("SMConv1d_v2") |
| 177 | # exit() |
| 178 | else: |
| 179 | return |
| 180 | |
| 181 | # help(m) |
| 182 | |
| 183 | # print("m is ", m) |
| 184 | |
| 185 | # print("std is", std) |
| 186 | |