| 593 | return out |
| 594 | |
| 595 | class Encoder(nn.Module): |
| 596 | def __init__(self, input_nc, output_nc, ngf=32, n_downsampling=4, norm_layer=nn.BatchNorm2d): |
| 597 | super(Encoder, self).__init__() |
| 598 | self.output_nc = output_nc |
| 599 | |
| 600 | model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), |
| 601 | norm_layer(ngf), nn.ReLU(True)] |
| 602 | ### downsample |
| 603 | for i in range(n_downsampling): |
| 604 | mult = 2**i |
| 605 | model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1), |
| 606 | norm_layer(ngf * mult * 2), nn.ReLU(True)] |
| 607 | |
| 608 | ### upsample |
| 609 | for i in range(n_downsampling): |
| 610 | mult = 2**(n_downsampling - i) |
| 611 | model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1), |
| 612 | norm_layer(int(ngf * mult / 2)), nn.ReLU(True)] |
| 613 | |
| 614 | model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()] |
| 615 | self.model = nn.Sequential(*model) |
| 616 | |
| 617 | def forward(self, input, inst): |
| 618 | outputs = self.model(input) |
| 619 | |
| 620 | # instance-wise average pooling |
| 621 | outputs_mean = outputs.clone() |
| 622 | for b in range(input.size()[0]): |
| 623 | inst_list = np.unique(inst[b].cpu().numpy().astype(int)) |
| 624 | for i in inst_list: |
| 625 | indices = (inst[b:b+1] == int(i)).nonzero() # n x 4 |
| 626 | for j in range(self.output_nc): |
| 627 | output_ins = outputs[indices[:,0] + b, indices[:,1] + j, indices[:,2], indices[:,3]] |
| 628 | mean_feat = torch.mean(output_ins).expand_as(output_ins) |
| 629 | ### add random noise to output feature |
| 630 | #mean_feat += torch.normal(torch.zeros_like(mean_feat), 0.05 * torch.ones_like(mean_feat)).cuda() |
| 631 | outputs_mean[indices[:,0] + b, indices[:,1] + j, indices[:,2], indices[:,3]] = mean_feat |
| 632 | return outputs_mean |
| 633 | |
| 634 | class MultiscaleDiscriminator(nn.Module): |
| 635 | def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, |