කටුසටහනක් RNN

32import math
33from typing import Optional, Tuple, Any
34
35import numpy as np
36import torch
37import torch.nn as nn
38from matplotlib import pyplot as plt
39from torch import optim
40from torch.utils.data import Dataset, DataLoader
41
42import einops
43from labml import lab, experiment, tracker, monit
44from labml_helpers.device import DeviceConfigs
45from labml_helpers.module import Module
46from labml_helpers.optimizer import OptimizerConfigs
47from labml_helpers.train_valid import TrainValidConfigs, hook_model_outputs, BatchIndex

50class StrokesDataset(Dataset):

57    def __init__(self, dataset: np.array, max_seq_length: int, scale: Optional[float] = None):

67        data = []

69        for seq in dataset:

71            if 10 < len(seq) <= max_seq_length:

73                seq = np.minimum(seq, 1000)
74                seq = np.maximum(seq, -1000)

76                seq = np.array(seq, dtype=np.float32)
77                data.append(seq)

83        if scale is None:
84            scale = np.std(np.concatenate([np.ravel(s[:, 0:2]) for s in data]))
85        self.scale = scale

88        longest_seq_len = max([len(seq) for seq in data])

98        self.data = torch.zeros(len(data), longest_seq_len + 2, 5, dtype=torch.float)

101        self.mask = torch.zeros(len(data), longest_seq_len + 1)
102
103        for i, seq in enumerate(data):
104            seq = torch.from_numpy(seq)
105            len_seq = len(seq)

107            self.data[i, 1:len_seq + 1, :2] = seq[:, :2] / scale

109            self.data[i, 1:len_seq + 1, 2] = 1 - seq[:, 2]

111            self.data[i, 1:len_seq + 1, 3] = seq[:, 2]

113            self.data[i, len_seq + 1:, 4] = 1

115            self.mask[i, :len_seq + 1] = 1

118        self.data[:, 0, 2] = 1

120    def __len__(self):

122        return len(self.data)

124    def __getitem__(self, idx: int):

126        return self.data[idx], self.mask[idx]

129class BivariateGaussianMixture:

139    def __init__(self, pi_logits: torch.Tensor, mu_x: torch.Tensor, mu_y: torch.Tensor,
140                 sigma_x: torch.Tensor, sigma_y: torch.Tensor, rho_xy: torch.Tensor):
141        self.pi_logits = pi_logits
142        self.mu_x = mu_x
143        self.mu_y = mu_y
144        self.sigma_x = sigma_x
145        self.sigma_y = sigma_y
146        self.rho_xy = rho_xy

148    @property
149    def n_distributions(self):

151        return self.pi_logits.shape[-1]

153    def set_temperature(self, temperature: float):

158        self.pi_logits /= temperature

160        self.sigma_x *= math.sqrt(temperature)

162        self.sigma_y *= math.sqrt(temperature)

164    def get_distribution(self):

166        sigma_x = torch.clamp_min(self.sigma_x, 1e-5)
167        sigma_y = torch.clamp_min(self.sigma_y, 1e-5)
168        rho_xy = torch.clamp(self.rho_xy, -1 + 1e-5, 1 - 1e-5)

171        mean = torch.stack([self.mu_x, self.mu_y], -1)

173        cov = torch.stack([
174            sigma_x * sigma_x, rho_xy * sigma_x * sigma_y,
175            rho_xy * sigma_x * sigma_y, sigma_y * sigma_y
176        ], -1)
177        cov = cov.view(*sigma_y.shape, 2, 2)

188        multi_dist = torch.distributions.MultivariateNormal(mean, covariance_matrix=cov)

191        cat_dist = torch.distributions.Categorical(logits=self.pi_logits)

194        return cat_dist, multi_dist

197class EncoderRNN(Module):

204    def __init__(self, d_z: int, enc_hidden_size: int):
205        super().__init__()

208        self.lstm = nn.LSTM(5, enc_hidden_size, bidirectional=True)

210        self.mu_head = nn.Linear(2 * enc_hidden_size, d_z)

212        self.sigma_head = nn.Linear(2 * enc_hidden_size, d_z)

214    def forward(self, inputs: torch.Tensor, state=None):

221        _, (hidden, cell) = self.lstm(inputs.float(), state)

225        hidden = einops.rearrange(hidden, 'fb b h -> b (fb h)')

228        mu = self.mu_head(hidden)

230        sigma_hat = self.sigma_head(hidden)

232        sigma = torch.exp(sigma_hat / 2.)

235        z = mu + sigma * torch.normal(mu.new_zeros(mu.shape), mu.new_ones(mu.shape))

238        return z, mu, sigma_hat

241class DecoderRNN(Module):

248    def __init__(self, d_z: int, dec_hidden_size: int, n_distributions: int):
249        super().__init__()

251        self.lstm = nn.LSTM(d_z + 5, dec_hidden_size)

255        self.init_state = nn.Linear(d_z, 2 * dec_hidden_size)

260        self.mixtures = nn.Linear(dec_hidden_size, 6 * n_distributions)

263        self.q_head = nn.Linear(dec_hidden_size, 3)

266        self.q_log_softmax = nn.LogSoftmax(-1)

269        self.n_distributions = n_distributions
270        self.dec_hidden_size = dec_hidden_size

272    def forward(self, x: torch.Tensor, z: torch.Tensor, state: Optional[Tuple[torch.Tensor, torch.Tensor]]):

274        if state is None:

276            h, c = torch.split(torch.tanh(self.init_state(z)), self.dec_hidden_size, 1)

279            state = (h.unsqueeze(0).contiguous(), c.unsqueeze(0).contiguous())

282        outputs, state = self.lstm(x, state)

285        q_logits = self.q_log_softmax(self.q_head(outputs))

291        pi_logits, mu_x, mu_y, sigma_x, sigma_y, rho_xy = \
292            torch.split(self.mixtures(outputs), self.n_distributions, 2)

305        dist = BivariateGaussianMixture(pi_logits, mu_x, mu_y,
306                                        torch.exp(sigma_x), torch.exp(sigma_y), torch.tanh(rho_xy))

309        return dist, q_logits, state

312class ReconstructionLoss(Module):

317    def forward(self, mask: torch.Tensor, target: torch.Tensor,
318                 dist: 'BivariateGaussianMixture', q_logits: torch.Tensor):

320        pi, mix = dist.get_distribution()

327        xy = target[:, :, 0:2].unsqueeze(-2).expand(-1, -1, dist.n_distributions, -1)

333        probs = torch.sum(pi.probs * torch.exp(mix.log_prob(xy)), 2)

342        loss_stroke = -torch.mean(mask * torch.log(1e-5 + probs))

345        loss_pen = -torch.mean(target[:, :, 2:] * q_logits)

348        return loss_stroke + loss_pen

351class KLDivLoss(Module):

358    def forward(self, sigma_hat: torch.Tensor, mu: torch.Tensor):

360        return -0.5 * torch.mean(1 + sigma_hat - mu ** 2 - torch.exp(sigma_hat))

363class Sampler:

370    def __init__(self, encoder: EncoderRNN, decoder: DecoderRNN):
371        self.decoder = decoder
372        self.encoder = encoder

374    def sample(self, data: torch.Tensor, temperature: float):

376        longest_seq_len = len(data)

379        z, _, _ = self.encoder(data)

382        s = data.new_tensor([0, 0, 1, 0, 0])
383        seq = [s]

386        state = None

389        with torch.no_grad():

391            for i in range(longest_seq_len):

393                data = torch.cat([s.view(1, 1, -1), z.unsqueeze(0)], 2)

396                dist, q_logits, state = self.decoder(data, z, state)

398                s = self._sample_step(dist, q_logits, temperature)

400                seq.append(s)

402                if s[4] == 1:
403                    break

406        seq = torch.stack(seq)

409        self.plot(seq)

411    @staticmethod
412    def _sample_step(dist: 'BivariateGaussianMixture', q_logits: torch.Tensor, temperature: float):

414        dist.set_temperature(temperature)

416        pi, mix = dist.get_distribution()

418        idx = pi.sample()[0, 0]

421        q = torch.distributions.Categorical(logits=q_logits / temperature)

423        q_idx = q.sample()[0, 0]

426        xy = mix.sample()[0, 0, idx]

429        stroke = q_logits.new_zeros(5)

431        stroke[:2] = xy

433        stroke[q_idx + 2] = 1

435        return stroke

437    @staticmethod
438    def plot(seq: torch.Tensor):

440        seq[:, 0:2] = torch.cumsum(seq[:, 0:2], dim=0)

442        seq[:, 2] = seq[:, 3]
443        seq = seq[:, 0:3].detach().cpu().numpy()

448        strokes = np.split(seq, np.where(seq[:, 2] > 0)[0] + 1)

450        for s in strokes:
451            plt.plot(s[:, 0], -s[:, 1])

453        plt.axis('off')

455        plt.show()

458class Configs(TrainValidConfigs):

466    device: torch.device = DeviceConfigs()

468    encoder: EncoderRNN
469    decoder: DecoderRNN
470    optimizer: optim.Adam
471    sampler: Sampler
472
473    dataset_name: str
474    train_loader: DataLoader
475    valid_loader: DataLoader
476    train_dataset: StrokesDataset
477    valid_dataset: StrokesDataset

480    enc_hidden_size = 256
481    dec_hidden_size = 512

484    batch_size = 100

487    d_z = 128

489    n_distributions = 20

492    kl_div_loss_weight = 0.5

494    grad_clip = 1.

496    temperature = 0.4

499    max_seq_length = 200
500
501    epochs = 100
502
503    kl_div_loss = KLDivLoss()
504    reconstruction_loss = ReconstructionLoss()

506    def init(self):

508        self.encoder = EncoderRNN(self.d_z, self.enc_hidden_size).to(self.device)
509        self.decoder = DecoderRNN(self.d_z, self.dec_hidden_size, self.n_distributions).to(self.device)

512        optimizer = OptimizerConfigs()
513        optimizer.parameters = list(self.encoder.parameters()) + list(self.decoder.parameters())
514        self.optimizer = optimizer

517        self.sampler = Sampler(self.encoder, self.decoder)

520        path = lab.get_data_path() / 'sketch' / f'{self.dataset_name}.npz'

522        dataset = np.load(str(path), encoding='latin1', allow_pickle=True)

525        self.train_dataset = StrokesDataset(dataset['train'], self.max_seq_length)

527        self.valid_dataset = StrokesDataset(dataset['valid'], self.max_seq_length, self.train_dataset.scale)

530        self.train_loader = DataLoader(self.train_dataset, self.batch_size, shuffle=True)

532        self.valid_loader = DataLoader(self.valid_dataset, self.batch_size)

535        hook_model_outputs(self.mode, self.encoder, 'encoder')
536        hook_model_outputs(self.mode, self.decoder, 'decoder')

539        tracker.set_scalar("loss.total.*", True)
540
541        self.state_modules = []

543    def step(self, batch: Any, batch_idx: BatchIndex):
544        self.encoder.train(self.mode.is_train)
545        self.decoder.train(self.mode.is_train)

550        data = batch[0].to(self.device).transpose(0, 1)
551        mask = batch[1].to(self.device).transpose(0, 1)

554        if self.mode.is_train:
555            tracker.add_global_step(len(data))

558        with monit.section("encoder"):

560            z, mu, sigma_hat = self.encoder(data)

563        with monit.section("decoder"):

565            z_stack = z.unsqueeze(0).expand(data.shape[0] - 1, -1, -1)
566            inputs = torch.cat([data[:-1], z_stack], 2)

568            dist, q_logits, _ = self.decoder(inputs, z, None)

571        with monit.section('loss'):

573            kl_loss = self.kl_div_loss(sigma_hat, mu)

575            reconstruction_loss = self.reconstruction_loss(mask, data[1:], dist, q_logits)

577            loss = reconstruction_loss + self.kl_div_loss_weight * kl_loss

580            tracker.add("loss.kl.", kl_loss)
581            tracker.add("loss.reconstruction.", reconstruction_loss)
582            tracker.add("loss.total.", loss)

585        if self.mode.is_train:

587            with monit.section('optimize'):

589                self.optimizer.zero_grad()

591                loss.backward()

593                if batch_idx.is_last:
594                    tracker.add(encoder=self.encoder, decoder=self.decoder)

596                nn.utils.clip_grad_norm_(self.encoder.parameters(), self.grad_clip)
597                nn.utils.clip_grad_norm_(self.decoder.parameters(), self.grad_clip)

599                self.optimizer.step()
600
601        tracker.save()

603    def sample(self):

605        data, *_ = self.valid_dataset[np.random.choice(len(self.valid_dataset))]

607        data = data.unsqueeze(1).to(self.device)

609        self.sampler.sample(data, self.temperature)

612def main():
613    configs = Configs()
614    experiment.create(name="sketch_rnn")

617    experiment.configs(configs, {
618        'optimizer.optimizer': 'Adam',

621        'optimizer.learning_rate': 1e-3,

623        'dataset_name': 'bicycle',

625        'inner_iterations': 10
626    })
627
628    with experiment.start():

630        configs.run()
631
632
633if __name__ == "__main__":
634    main()