"""AlignDRAW 512x512 Based on the supplied AlignDRAW PyTorch implementation. Image resolution: 512x512 DRAW generation steps: 69 Requires a newly trained checkpoint. """ import torch import torch.nn as nn import torch.nn.functional as F def _lstm(pre, cell, H): i = torch.sigmoid(pre[:, 0:H]) f = torch.sigmoid(pre[:, H:2*H]) g = torch.tanh(pre[:, 2*H:3*H]) o = torch.sigmoid(pre[:, 3*H:4*H]) cell = f * cell + i * g return o * torch.tanh(cell), cell class SelectiveAttention: def __init__(self, H, W, N): self.H, self.W, self.N = H, W, N def params(self, raw): gx = (self.W + 1) / 2.0 * (raw[:, 0] + 1) gy = (self.H + 1) / 2.0 * (raw[:, 1] + 1) delta = (max(self.H, self.W) - 1.0) / (self.N - 1) * torch.exp(raw[:, 2]) sigma = torch.exp(raw[:, 3] / 2.0) gamma = torch.exp(raw[:, 4]).unsqueeze(1) return gx, gy, delta, sigma, gamma def _fb(self, g, delta, sigma, size): grid = torch.arange( self.N, device=g.device, dtype=g.dtype ) - self.N // 2 - 0.5 mu = g.unsqueeze(1) + delta.unsqueeze(1) * grid.unsqueeze(0) a = torch.arange(size, device=g.device, dtype=g.dtype) fb = torch.exp( -((a.view(1, 1, -1) - mu.unsqueeze(2)) ** 2) / (2 * sigma.view(-1, 1, 1) ** 2 + 1e-8) ) return fb / (fb.sum(-1, keepdim=True) + 1e-4) def read(self, img, gx, gy, delta, sigma): Fx = self._fb(gx, delta, sigma, self.W) Fy = self._fb(gy, delta, sigma, self.H) t = torch.einsum('bnh,bchw->bcnw', Fy, img) return torch.einsum('bcnw,bmw->bcnm', t, Fx) def write(self, win, gx, gy, delta, sigma): Fx = self._fb(gx, delta, sigma, self.W) Fy = self._fb(gy, delta, sigma, self.H) t = torch.einsum('bnh,bcnm->bchm', Fy, win) return torch.einsum('bchm,bmw->bchw', t, Fx) class AlignDRAW(nn.Module): def __init__( self, dimY=25323, dimLangRNN=128, dimAlign=512, imgSize=512, dimReadAttent=9, dimWriteAttent=9, dimRNNEnc=550, dimRNNDec=550, dimZ=275, runSteps=69 ): super().__init__() self.C = 3 self.H = self.W = imgSize self.dimX = self.C * self.H * self.W self.Nr, self.Nw = dimReadAttent, dimWriteAttent self.He, self.Hd = dimRNNEnc, dimRNNDec self.Z, self.T = dimZ, runSteps self.Hl, self.Al = dimLangRNN, dimAlign # Bidirectional language encoder self.emb_f = nn.Embedding(dimY, 4 * dimLangRNN) self.Wh_f = nn.Linear(dimLangRNN, 4 * dimLangRNN, bias=False) self.b_f = nn.Parameter(torch.zeros(4 * dimLangRNN)) self.emb_b = nn.Embedding(dimY, 4 * dimLangRNN) self.Wh_b = nn.Linear(dimLangRNN, 4 * dimLangRNN, bias=False) self.b_b = nn.Parameter(torch.zeros(4 * dimLangRNN)) # Caption alignment self.W_lang_align = nn.Linear( 2 * dimLangRNN, dimAlign, bias=False ) self.W_hdec_align = nn.Linear( dimRNNDec, dimAlign, bias=False ) self.b_align = nn.Parameter(torch.zeros(dimAlign)) self.v_align = nn.Parameter(torch.randn(dimAlign) * 0.01) self.W_s_hdec = nn.Linear( 2 * dimLangRNN, 4 * dimRNNDec, bias=False ) # Image reader self.W_read = nn.Linear(dimRNNDec, 5) read_dim = 2 * self.C * self.Nr * self.Nr # Encoder LSTM self.Wh_enc = nn.Linear( dimRNNEnc, 4 * dimRNNEnc, bias=False ) self.Wi_enc = nn.Linear( read_dim + dimRNNDec, 4 * dimRNNEnc, bias=False ) self.b_enc = nn.Parameter(torch.zeros(4 * dimRNNEnc)) self.W_mu = nn.Linear(dimRNNEnc, dimZ) self.W_logsigma = nn.Linear(dimRNNEnc, dimZ) # Decoder and learned prior self.Wh_dec = nn.Linear( dimRNNDec, 4 * dimRNNDec, bias=False ) self.Wz_dec = nn.Linear( dimZ, 4 * dimRNNDec, bias=False ) self.b_dec = nn.Parameter(torch.zeros(4 * dimRNNDec)) self.W_prior = nn.Linear(dimRNNDec, 2 * dimZ) # Image writer self.W_write = nn.Linear(dimRNNDec, 5) self.W_c = nn.Linear( dimRNNDec, self.C * self.Nw * self.Nw ) # Full 512x512 canvas initialization self.h0_lang = nn.Parameter(torch.zeros(1, dimLangRNN)) self.h0_enc = nn.Parameter(torch.zeros(1, dimRNNEnc)) self.h0_dec = nn.Parameter(torch.zeros(1, dimRNNDec)) self.c0 = nn.Parameter(torch.zeros(1, self.dimX)) self.reader = SelectiveAttention( self.H, self.W, self.Nr ) self.writer = SelectiveAttention( self.H, self.W, self.Nw ) self._init_weights() def _init_weights(self): for m in self.modules(): if isinstance(m, (nn.Linear, nn.Embedding)): nn.init.normal_(m.weight, 0.0, 0.01) if getattr(m, "bias", None) is not None: nn.init.zeros_(m.bias) def encode_caption(self, y): b, L = y.shape yr = torch.flip(y, dims=[1]) hf = self.h0_lang.expand(b, -1).contiguous() hb = self.h0_lang.expand(b, -1).contiguous() cf = y.new_zeros((b, self.Hl), dtype=torch.float) cb = y.new_zeros((b, self.Hl), dtype=torch.float) outs_f, outs_b = [], [] for t in range(L): hf, cf = _lstm( self.emb_f(y[:, t]) + self.Wh_f(hf) + self.b_f, cf, self.Hl ) hb, cb = _lstm( self.emb_b(yr[:, t]) + self.Wh_b(hb) + self.b_b, cb, self.Hl ) outs_f.append(hf) outs_b.append(hb) hlang = torch.cat( [ torch.stack(outs_f, 1), torch.stack(outs_b, 1) ], dim=2 ) return hlang, self.W_lang_align(hlang) def _align(self, hlang, hid_align, h_dec): e = torch.tanh( hid_align + self.W_hdec_align(h_dec).unsqueeze(1) + self.b_align ) e = (e * self.v_align).sum(-1) alpha = F.softmax(e, dim=1) s = (alpha.unsqueeze(2) * hlang).sum(1) return s, alpha def forward(self, x, y): b = x.shape[0] assert x.shape[1] == self.dimX, ( f"Expected {self.dimX} image values, got {x.shape[1]}" ) h_dec = self.h0_dec.expand(b, -1).contiguous() cell_dec = x.new_zeros(b, self.Hd) h_enc = self.h0_enc.expand(b, -1).contiguous() cell_enc = x.new_zeros(b, self.He) c = self.c0.expand(b, -1).contiguous() mu_p = x.new_zeros(b, self.Z) ls_p = x.new_zeros(b, self.Z) hlang, hid_align = self.encode_caption(y) ximg = x.view(b, self.C, self.H, self.W) kl = x.new_zeros(()) for _ in range(self.T): x_hat = ximg - torch.sigmoid(c).view( b, self.C, self.H, self.W ) gx, gy, dl, sg, gm = self.reader.params( self.W_read(h_dec) ) r = gm[:, :, None, None] * torch.cat( [ self.reader.read(ximg, gx, gy, dl, sg), self.reader.read(x_hat, gx, gy, dl, sg) ], dim=1 ) r = r.reshape(b, -1) s, _ = self._align(hlang, hid_align, h_dec) h_enc, cell_enc = _lstm( self.Wh_enc(h_enc) + self.Wi_enc(torch.cat([r, h_dec], 1)) + self.b_enc, cell_enc, self.He ) mu = self.W_mu(h_enc) ls = 0.5 * self.W_logsigma(h_enc) z = mu + torch.exp(ls) * torch.randn_like(mu) kl = kl + torch.sum( -1 + ( (mu - mu_p) ** 2 + torch.exp(2 * ls) ) / torch.exp(2 * ls_p) - 2 * ls + 2 * ls_p ) h_dec, cell_dec = _lstm( self.Wh_dec(h_dec) + self.Wz_dec(z) + self.W_s_hdec(s) + self.b_dec, cell_dec, self.Hd ) mls = torch.tanh(self.W_prior(h_dec)) mu_p = mls[:, :self.Z] ls_p = mls[:, self.Z:] gx, gy, dl, sg, gm = self.writer.params( self.W_write(h_dec) ) win = self.W_c(h_dec).view( b, self.C, self.Nw, self.Nw ) c = c + ( 1.0 / gm[:, :, None, None] * self.writer.write(win, gx, gy, dl, sg) ).reshape(b, -1) c_final = torch.sigmoid(c) bce = F.binary_cross_entropy( c_final, x, reduction="sum" ) kl = 0.5 * kl return { "loss": (bce + kl) / b, "bce": bce / b, "kl": kl / b } @torch.no_grad() def generate(self, y, steps=None, return_seq=False): steps = self.T if steps is None else steps b = y.shape[0] dev = self.h0_dec.device h_dec = self.h0_dec.expand(b, -1).contiguous() cell_dec = torch.zeros( b, self.Hd, device=dev ) c = self.c0.expand(b, -1).contiguous() mu_p = torch.zeros(b, self.Z, device=dev) ls_p = torch.zeros(b, self.Z, device=dev) hlang, hid_align = self.encode_caption(y) seq = [] for _ in range(steps): z = mu_p + torch.exp(ls_p) * torch.randn_like(mu_p) s, _ = self._align( hlang, hid_align, h_dec ) h_dec, cell_dec = _lstm( self.Wh_dec(h_dec) + self.Wz_dec(z) + self.W_s_hdec(s) + self.b_dec, cell_dec, self.Hd ) mls = torch.tanh(self.W_prior(h_dec)) mu_p = mls[:, :self.Z] ls_p = mls[:, self.Z:] gx, gy, dl, sg, gm = self.writer.params( self.W_write(h_dec) ) win = self.W_c(h_dec).view( b, self.C, self.Nw, self.Nw ) c = c + ( 1.0 / gm[:, :, None, None] * self.writer.write(win, gx, gy, dl, sg) ).reshape(b, -1) if return_seq: seq.append( torch.sigmoid(c).view( b, self.C, self.H, self.W ).cpu() ) img = torch.sigmoid(c).view( b, self.C, self.H, self.W ) return (img, seq) if return_seq else img