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 __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 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 make_decoder(self):
filt_dims = self.filt_dims[::-1]
self.dec_res = ResidualBlock(2 * filt_dims[-1], filt_dims[-1])
dec = nn.ModuleList()
for i_dim in range(len(filt_dims) - 1):
dec.append([
TransposeConvBlock(filt_dims[i_dim],
TransposeConvBlock[i_dim + 1])
])
self.dec = dec
self.dec_im_size = self.im_size * 2**len(dec)
def build_decoder(self):
conv_repeat = 3
resnet_repeat = 2
cur_filts = self.encoded_input_filts
decoder_block = []
for _ in range(resnet_repeat):
decoder_block.append(ResidualBlock(cur_filts, cur_filts))
for _ in range(conv_repeat):
decoder_block.append(TransposeConvBlock(cur_filts, cur_filts // 2))
cur_filts = cur_filts // 2
self.decoder = nn.Sequential(*decoder_block)
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)