File size: 2,714 Bytes
18dd6ad |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
import abc
from typing import Tuple, List
import torch
import torch.nn as nn
from annotator.lama.saicinpainting.training.modules.depthwise_sep_conv import DepthWiseSeperableConv
from annotator.lama.saicinpainting.training.modules.multidilated_conv import MultidilatedConv
class BaseDiscriminator(nn.Module):
@abc.abstractmethod
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
"""
Predict scores and get intermediate activations. Useful for feature matching loss
:return tuple (scores, list of intermediate activations)
"""
raise NotImplemented()
def get_conv_block_ctor(kind='default'):
if not isinstance(kind, str):
return kind
if kind == 'default':
return nn.Conv2d
if kind == 'depthwise':
return DepthWiseSeperableConv
if kind == 'multidilated':
return MultidilatedConv
raise ValueError(f'Unknown convolutional block kind {kind}')
def get_norm_layer(kind='bn'):
if not isinstance(kind, str):
return kind
if kind == 'bn':
return nn.BatchNorm2d
if kind == 'in':
return nn.InstanceNorm2d
raise ValueError(f'Unknown norm block kind {kind}')
def get_activation(kind='tanh'):
if kind == 'tanh':
return nn.Tanh()
if kind == 'sigmoid':
return nn.Sigmoid()
if kind is False:
return nn.Identity()
raise ValueError(f'Unknown activation kind {kind}')
class SimpleMultiStepGenerator(nn.Module):
def __init__(self, steps: List[nn.Module]):
super().__init__()
self.steps = nn.ModuleList(steps)
def forward(self, x):
cur_in = x
outs = []
for step in self.steps:
cur_out = step(cur_in)
outs.append(cur_out)
cur_in = torch.cat((cur_in, cur_out), dim=1)
return torch.cat(outs[::-1], dim=1)
def deconv_factory(kind, ngf, mult, norm_layer, activation, max_features):
if kind == 'convtranspose':
return [nn.ConvTranspose2d(min(max_features, ngf * mult),
min(max_features, int(ngf * mult / 2)),
kernel_size=3, stride=2, padding=1, output_padding=1),
norm_layer(min(max_features, int(ngf * mult / 2))), activation]
elif kind == 'bilinear':
return [nn.Upsample(scale_factor=2, mode='bilinear'),
DepthWiseSeperableConv(min(max_features, ngf * mult),
min(max_features, int(ngf * mult / 2)),
kernel_size=3, stride=1, padding=1),
norm_layer(min(max_features, int(ngf * mult / 2))), activation]
else:
raise Exception(f"Invalid deconv kind: {kind}") |