def __init__(self, im_size, num_classes):
super().__init__()
# feed_input = 10
feed_input = 8
q = int(np.log2(im_size))
filt_dims = [2**(12 - i) for i in range(q, 1, -1)]
print("Discriminator Filters: ", filt_dims)
self.from_rgb = conv311(feed_input, filt_dims[0])
block = []
for i in range(len(filt_dims) - 1):
block.append(ConvBlock(filt_dims[i], filt_dims[i + 1]))
if i != len(filt_dims) - 2:
block.append(ResidualBlock(filt_dims[i + 1], filt_dims[i + 1]))
block.append(ConvBlock(filt_dims[-1], filt_dims[-1]))
self.down_sample = nn.Sequential(*block)
self.conv_real = nn.Conv2d(filt_dims[-1],
1,
kernel_size=2,
stride=1,
padding=0)
self.conv_class = nn.Conv2d(filt_dims[-1],
num_classes,
kernel_size=2,
stride=1,
padding=0)
self.num_classes = num_classes
def __init__(self, im_size):
super().__init__()
self.im_size = im_size
# feed to encoder
# feed_input = 7
# feed_input = 9
feed_input = 5
filt_dims = [32, 64, 128, 256]
filt_const = 1
self.to_enc = conv311(feed_input, 32 * filt_const)
self.e1 = ConvBlock(32 * filt_const, 64 * filt_const)
self.e2 = ConvBlock(64 * filt_const, 128 * filt_const)
self.e2_e3 = SelfAttnBlock(128 * filt_const)
self.e3 = ConvBlock(128 * filt_const, 256 * filt_const)
self.res = nn.Sequential(
ResidualBlock(256 * filt_const, 256 * filt_const),
ResidualBlock(256 * filt_const, 256 * filt_const))
self.d3 = TransposeConvBlock(256 * filt_const, 128 * filt_const)
self.d3_d2 = SelfAttnBlock(128 * filt_const)
self.d2 = TransposeConvBlock(256 * filt_const, 64 * filt_const)
self.d1 = TransposeConvBlock(64 * filt_const, 32 * filt_const)
self.from_dec = conv311(32 * filt_const, 3)
def __init__(self, im_size):
super().__init__()
self.im_size = im_size
feed_input = 5
filt_dims = [32, 64, 128, 256]
filt_const = 1
self.to_enc = conv311(feed_input, 32 * filt_const)
self.enc = nn.Sequential(ConvBlock(32 * filt_const, 64 * filt_const),
ConvBlock(64 * filt_const, 128 * filt_const),
SelfAttnBlock(128 * filt_const),
ConvBlock(128 * filt_const, 256 * filt_const))
self.res = nn.Sequential(
ResidualBlock(256 * filt_const, 256 * filt_const),
ResidualBlock(256 * filt_const, 256 * filt_const))
self.dec = nn.Sequential(
TransposeConvBlock(256 * filt_const, 128 * filt_const),
SelfAttnBlock(128 * filt_const),
TransposeConvBlock(128 * filt_const, 64 * filt_const),
# SelfAttnBlock(64 * filt_const)
TransposeConvBlock(64 * filt_const, 32 * filt_const),
)
self.from_dec = conv311(32 * filt_const, 3)
def build_encoder_input(self):
filt_const = self.filt_const
self.enc_input = nn.Sequential(
conv311(4, 32 * self.filt_const),
ConvBlock(32 * filt_const, 64 * filt_const),
ConvBlock(64 * filt_const, 128 * filt_const),
SelfAttnBlock(128 * filt_const),
ConvBlock(128 * filt_const, 256 * filt_const))
self.encoded_input_filts = 256 * filt_const
def build_encoder_template(self):
template_dim = 1
conv_filts = 4
encoder_block = [conv311(template_dim, conv_filts)]
for _ in range(conv_repeat):
encoder_block.append(ConvBlock(conv_filts, conv_filts))
self.encode_template = nn.Sequential(*encoder_block)
self.encoded_template_filts = conv_filts
def make_encoder(self):
filt_dims = self.filt_dims
enc = nn.ModuleList()
for i_dim in range(len(filt_dims) - 1):
enc.append([ConvBlock(filt_dims[i_dim], filt_dims[i_dim + 1])])
self.enc = enc
self.enc_res = ResidualBlock(filt_dims[-1], filt_dims[-1])
self.enc_im_size = self.im_size // 2**len(enc)
def __init__(self, dev_id):
super().__init__()
self.dev_id = dev_id
self.noise_dim = 50
feed_input = 5
filt_dims = [32, 64, 128, 256]
filt_const = 2
self.to_enc = conv311(feed_input, 32 * filt_const)
self.e1 = ConvBlock(32 * filt_const, 64 * filt_const)
self.e2 = ConvBlock(64 * filt_const, 128 * filt_const)
self.e3 = ConvBlock(128 * filt_const, 256 * filt_const)
# self.e3 = ConvBlock(128, 128) for attention change to 128
self.res = nn.Sequential(
ResidualBlock(256 * filt_const, 256 * filt_const),
ResidualBlock(256 * filt_const, 256 * filt_const))
self.d3 = TransposeConvBlock(256 * filt_const + self.noise_dim // 2,
128 * filt_const)
self.d2 = TransposeConvBlock(256 * filt_const, 64 * filt_const)
self.d1 = TransposeConvBlock(128 * filt_const, 32 * filt_const)
self.from_dec = conv311(32 * filt_const, 3)
self.noise_feat = nn.Sequential(
TransposeConvBlock(self.noise_dim,
self.noise_dim // 2,
k=2,
s=1,
p=1),
TransposeConvBlock(
self.noise_dim // 2,
self.noise_dim // 2,
k=2,
), TransposeConvBlock(self.noise_dim // 2, self.noise_dim // 2),
TransposeConvBlock(self.noise_dim // 2, self.noise_dim // 2))
def build_encoder_input(self):
conv_filters = self.conv_filters
conv_repeat = self.conv_repeat
resnet_repeat = self.resnet_repeat
encoder_block = []
cur_filts = conv_filters
for _ in range(conv_repeat):
encoder_block.append(ConvBlock(cur_filts, cur_filts * 2))
cur_filts = cur_filts * 2
for _ in range(resnet_repeat):
encoder_block.append(ResidualBlock(cur_filts, cur_filts))
self.encoded_input_filts = cur_filts
self.encode_input = nn.Sequential(*encoder_block)
def __init__(self, image_size):
super().__init__()
feed_input = 5 + 4
conv_dim = 32
self.layer_1 = nn.Sequential(
nn.Conv2d(feed_input, conv_dim, kernel_size=3, stride=1,
padding=1), ConvBlock(conv_dim, conv_dim, k=3, s=1))
self.layer_2 = nn.Sequential(
ConvBlock(conv_dim, conv_dim * 2),
ConvBlock(conv_dim * 2, conv_dim * 2, k=3, s=1))
self.layer_3 = nn.Sequential(
ConvBlock(conv_dim * 2, conv_dim * 4),
ConvBlock(conv_dim * 4, conv_dim * 4, k=3, s=1))
self.layer_4 = nn.Sequential(
ConvBlock(conv_dim * 4, conv_dim * 8),
ConvBlock(conv_dim * 8, conv_dim * 8, k=3, s=1))
self.bottom_upsample = nn.Sequential(
ConvBlock(conv_dim * 8, conv_dim * 16),
ConvBlock(conv_dim * 16, conv_dim * 32, k=3, s=1, p=1),
TransposeConvBlock(conv_dim * 32, conv_dim * 8)) # 8 - > 16
self.upsampling_1 = nn.Sequential(
ConvBlock(conv_dim * 16, conv_dim * 8, k=3, s=1, p=1),
ConvBlock(conv_dim * 8, conv_dim * 8, k=3, s=1),
TransposeConvBlock(conv_dim * 8, conv_dim * 4)) # 16 -> 32
self.upsampling_2 = nn.Sequential(
ConvBlock(conv_dim * 8, conv_dim * 4, k=3, s=1, p=1),
ConvBlock(conv_dim * 4, conv_dim * 4, k=3, s=1),
TransposeConvBlock(conv_dim * 4, conv_dim * 2)) # 32 -> 64
self.upsampling_3 = nn.Sequential(
ConvBlock(conv_dim * 4, conv_dim * 2, k=3, s=1, p=1),
ConvBlock(conv_dim * 2, conv_dim * 2, k=3, s=1),
TransposeConvBlock(conv_dim * 2, conv_dim)) # 64 -> 128
self.final = nn.Sequential(
ConvBlock(conv_dim * 2, conv_dim, k=3, s=1, p=1),
nn.Conv2d(conv_dim, 3, stride=1, padding=1, kernel_size=3))