| 518 | |
| 519 | |
| 520 | class StridedConvF(nn.Module): |
| 521 | def __init__(self, init_type='normal', init_gain=0.02, gpu_ids=[]): |
| 522 | super().__init__() |
| 523 | # self.conv1 = nn.Conv2d(256, 128, 3, stride=2) |
| 524 | # self.conv2 = nn.Conv2d(128, 64, 3, stride=1) |
| 525 | self.l2_norm = Normalize(2) |
| 526 | self.mlps = {} |
| 527 | self.moving_averages = {} |
| 528 | self.init_type = init_type |
| 529 | self.init_gain = init_gain |
| 530 | self.gpu_ids = gpu_ids |
| 531 | |
| 532 | def create_mlp(self, x): |
| 533 | C, H = x.shape[1], x.shape[2] |
| 534 | n_down = int(np.rint(np.log2(H / 32))) |
| 535 | mlp = [] |
| 536 | for i in range(n_down): |
| 537 | mlp.append(nn.Conv2d(C, max(C // 2, 64), 3, stride=2)) |
| 538 | mlp.append(nn.ReLU()) |
| 539 | C = max(C // 2, 64) |
| 540 | mlp.append(nn.Conv2d(C, 64, 3)) |
| 541 | mlp = nn.Sequential(*mlp) |
| 542 | init_net(mlp, self.init_type, self.init_gain, self.gpu_ids) |
| 543 | return mlp |
| 544 | |
| 545 | def update_moving_average(self, key, x): |
| 546 | if key not in self.moving_averages: |
| 547 | self.moving_averages[key] = x.detach() |
| 548 | |
| 549 | self.moving_averages[key] = self.moving_averages[key] * 0.999 + x.detach() * 0.001 |
| 550 | |
| 551 | def forward(self, x, use_instance_norm=False): |
| 552 | C, H = x.shape[1], x.shape[2] |
| 553 | key = '%d_%d' % (C, H) |
| 554 | if key not in self.mlps: |
| 555 | self.mlps[key] = self.create_mlp(x) |
| 556 | self.add_module("child_%s" % key, self.mlps[key]) |
| 557 | mlp = self.mlps[key] |
| 558 | x = mlp(x) |
| 559 | self.update_moving_average(key, x) |
| 560 | x = x - self.moving_averages[key] |
| 561 | if use_instance_norm: |
| 562 | x = F.instance_norm(x) |
| 563 | return self.l2_norm(x) |
| 564 | |
| 565 | |
| 566 | class PatchSampleF(nn.Module): |