attention/pipelines.py at main · Xylambda/attention · GitHub

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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
import torch
import numpy as np
from utils import Pipeline
from models import (
    AttentionLSTM,
    VanillaLSTM,
    EmbeddingLSTM,
    AttentionEmbeddingLSTM
)

torch.manual_seed(0)
np.random.seed(0)


class AttentionLSTMPipeline(Pipeline):
    def __init__(self, dataset="sine_wave"):
        super().__init__()

        self.dataset = dataset

    def create_model(self):
        self.generate_data()

        X_train = self.X_train
        y_train = self.y_train

        model = AttentionLSTM(
            embed_dim=X_train.shape[2], out_size=y_train.shape[-1]
        )
        self.model = model


class VanillaLSTMPipeline(Pipeline):
    def __init__(self, dataset="sine_wave"):
        super().__init__()

        self.dataset = dataset

    def create_model(self):
        self.generate_data()

        X_train = self.X_train
        y_train = self.y_train

        model = VanillaLSTM(
            input_size=X_train.shape[2], out_size=y_train.shape[-1]
        )
        self.model = model


class EmbeddingLSTMPipeline(Pipeline):
    def __init__(self, dataset="sine_wave"):
        super().__init__()

        self.dataset = dataset

    def create_model(self):
        self.generate_data()

        X_train = self.X_train
        y_train = self.y_train
        features = X_train.shape[1]
        mini_batch = X_train.shape[2]

        model = EmbeddingLSTM(
            linear_channel=features,
            period_channel=(mini_batch - features),
            input_channel=mini_batch,
            input_size=X_train.shape[2],
            out_size=y_train.shape[-1]
        )
        self.model = model


class AttentionEmbeddingLSTMPipeline(Pipeline):
    def __init__(self, dataset="sine_wave"):
        super().__init__()

        self.dataset = dataset

    def create_model(self):
        self.generate_data()

        X_train = self.X_train
        y_train = self.y_train
        features = X_train.shape[1]
        mini_batch = X_train.shape[2]

        model = AttentionEmbeddingLSTM(
            linear_channel=features,
            period_channel=(mini_batch - features),
            input_channel=mini_batch,
            input_size=X_train.shape[2],
            out_size=y_train.shape[-1]
        )
        self.model = model