--eval
35 | ```
36 | The detailed descriptions about the arguments are as following:
37 | | Parameter name | Description of parameter |
38 | | --- | --- |
39 | | dataset_name | The dataset name |
40 | | run_name | The folder name used to save model, output and evaluation metrics. This can be set to any word |
41 | | archive | The archive name that the dataset belongs to. This can be set to `forecast_csv` or `forecast_csv_univar` |
42 | | batch_size | The batch size (defaults to 8) |
43 | | repr_dims | The representation dimensions (defaults to 320) |
44 | | gpu | The gpu no. used for training and inference (defaults to 0) |
45 | | eval | Whether to perform evaluation after training |
46 | | kernels | Kernel sizes for mixture of AR experts module |
47 | | alpha | Weight for loss function |
48 |
49 | (For descriptions of more arguments, run `python train.py -h`.)
50 |
51 | ## Main Results
52 | We perform experiments on five real-world public benchmark datasets, comparing against both state-of-the-art representation learning and end-to-end forecasting approaches.
53 | CoST achieves state-of-the-art performance, beating the best performing end-to-end forecasting approach by 39.3% and 18.22% (MSE) in the multivariate and univariate settings
54 | respectively. CoST also beats next best performing feature-based approach by 21.3% and 4.71% (MSE) in the multivariate and univariate settings respectively (refer to main paper for full results).
55 |
56 |
57 | 
58 |
59 |
60 | ## FAQs
61 | **Q**: ValueError: Found array with dim 4. StandardScaler expected <= 2.
62 |
63 | **A**: Please install the appropriate package requirements as found in ```requirements.txt```, in particular, ```scikit_learn==0.24.1```.
64 |
65 | **Q**: How to set the ``--kernels`` parameter?
66 |
67 | **A**: It should be list of space separated integers, e.g. ```--kernels 1 2 4```. See the `scripts` folder for further examples.
68 |
69 | ## Acknowledgements
70 | The implementation of CoST relies on resources from the following codebases and repositories, we thank the original authors for open-sourcing their work.
71 | * https://github.com/yuezhihan/ts2vec
72 | * https://github.com/zhouhaoyi/Informer2020
73 |
74 | ## Citation
75 | Please consider citing if you find this code useful to your research.
76 | @inproceedings{
77 | woo2022cost,
78 | title={Co{ST}: Contrastive Learning of Disentangled Seasonal-Trend Representations for Time Series Forecasting},
79 | author={Gerald Woo and Chenghao Liu and Doyen Sahoo and Akshat Kumar and Steven Hoi},
80 | booktitle={International Conference on Learning Representations},
81 | year={2022},
82 | url={https://openreview.net/forum?id=PilZY3omXV2}
83 | }
84 |
--------------------------------------------------------------------------------
/CODE_OF_CONDUCT.md:
--------------------------------------------------------------------------------
1 | # Salesforce Open Source Community Code of Conduct
2 |
3 | ## About the Code of Conduct
4 |
5 | Equality is a core value at Salesforce. We believe a diverse and inclusive
6 | community fosters innovation and creativity, and are committed to building a
7 | culture where everyone feels included.
8 |
9 | Salesforce open-source projects are committed to providing a friendly, safe, and
10 | welcoming environment for all, regardless of gender identity and expression,
11 | sexual orientation, disability, physical appearance, body size, ethnicity, nationality,
12 | race, age, religion, level of experience, education, socioeconomic status, or
13 | other similar personal characteristics.
14 |
15 | The goal of this code of conduct is to specify a baseline standard of behavior so
16 | that people with different social values and communication styles can work
17 | together effectively, productively, and respectfully in our open source community.
18 | It also establishes a mechanism for reporting issues and resolving conflicts.
19 |
20 | All questions and reports of abusive, harassing, or otherwise unacceptable behavior
21 | in a Salesforce open-source project may be reported by contacting the Salesforce
22 | Open Source Conduct Committee at ossconduct@salesforce.com.
23 |
24 | ## Our Pledge
25 |
26 | In the interest of fostering an open and welcoming environment, we as
27 | contributors and maintainers pledge to making participation in our project and
28 | our community a harassment-free experience for everyone, regardless of gender
29 | identity and expression, sexual orientation, disability, physical appearance,
30 | body size, ethnicity, nationality, race, age, religion, level of experience, education,
31 | socioeconomic status, or other similar personal characteristics.
32 |
33 | ## Our Standards
34 |
35 | Examples of behavior that contributes to creating a positive environment
36 | include:
37 |
38 | * Using welcoming and inclusive language
39 | * Being respectful of differing viewpoints and experiences
40 | * Gracefully accepting constructive criticism
41 | * Focusing on what is best for the community
42 | * Showing empathy toward other community members
43 |
44 | Examples of unacceptable behavior by participants include:
45 |
46 | * The use of sexualized language or imagery and unwelcome sexual attention or
47 | advances
48 | * Personal attacks, insulting/derogatory comments, or trolling
49 | * Public or private harassment
50 | * Publishing, or threatening to publish, others' private information—such as
51 | a physical or electronic address—without explicit permission
52 | * Other conduct which could reasonably be considered inappropriate in a
53 | professional setting
54 | * Advocating for or encouraging any of the above behaviors
55 |
56 | ## Our Responsibilities
57 |
58 | Project maintainers are responsible for clarifying the standards of acceptable
59 | behavior and are expected to take appropriate and fair corrective action in
60 | response to any instances of unacceptable behavior.
61 |
62 | Project maintainers have the right and responsibility to remove, edit, or
63 | reject comments, commits, code, wiki edits, issues, and other contributions
64 | that are not aligned with this Code of Conduct, or to ban temporarily or
65 | permanently any contributor for other behaviors that they deem inappropriate,
66 | threatening, offensive, or harmful.
67 |
68 | ## Scope
69 |
70 | This Code of Conduct applies both within project spaces and in public spaces
71 | when an individual is representing the project or its community. Examples of
72 | representing a project or community include using an official project email
73 | address, posting via an official social media account, or acting as an appointed
74 | representative at an online or offline event. Representation of a project may be
75 | further defined and clarified by project maintainers.
76 |
77 | ## Enforcement
78 |
79 | Instances of abusive, harassing, or otherwise unacceptable behavior may be
80 | reported by contacting the Salesforce Open Source Conduct Committee
81 | at ossconduct@salesforce.com. All complaints will be reviewed and investigated
82 | and will result in a response that is deemed necessary and appropriate to the
83 | circumstances. The committee is obligated to maintain confidentiality with
84 | regard to the reporter of an incident. Further details of specific enforcement
85 | policies may be posted separately.
86 |
87 | Project maintainers who do not follow or enforce the Code of Conduct in good
88 | faith may face temporary or permanent repercussions as determined by other
89 | members of the project's leadership and the Salesforce Open Source Conduct
90 | Committee.
91 |
92 | ## Attribution
93 |
94 | This Code of Conduct is adapted from the [Contributor Covenant][contributor-covenant-home],
95 | version 1.4, available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html.
96 | It includes adaptions and additions from [Go Community Code of Conduct][golang-coc],
97 | [CNCF Code of Conduct][cncf-coc], and [Microsoft Open Source Code of Conduct][microsoft-coc].
98 |
99 | This Code of Conduct is licensed under the [Creative Commons Attribution 3.0 License][cc-by-3-us].
100 |
101 | [contributor-covenant-home]: https://www.contributor-covenant.org (https://www.contributor-covenant.org/)
102 | [golang-coc]: https://golang.org/conduct
103 | [cncf-coc]: https://github.com/cncf/foundation/blob/master/code-of-conduct.md
104 | [microsoft-coc]: https://opensource.microsoft.com/codeofconduct/
105 | [cc-by-3-us]: https://creativecommons.org/licenses/by/3.0/us/
--------------------------------------------------------------------------------
/train.py:
--------------------------------------------------------------------------------
1 | import argparse
2 | import os
3 | import time
4 | import datetime
5 | import math
6 | import numpy as np
7 | import tasks
8 | import datautils
9 | from utils import init_dl_program, name_with_datetime, pkl_save, data_dropout
10 |
11 | # import methods
12 | from cost import CoST
13 |
14 |
15 | def save_checkpoint_callback(
16 | save_every=1,
17 | unit='epoch'
18 | ):
19 | assert unit in ('epoch', 'iter')
20 | def callback(model, loss):
21 | n = model.n_epochs if unit == 'epoch' else model.n_iters
22 | if n % save_every == 0:
23 | model.save(f'{run_dir}/model_{n}.pkl')
24 | return callback
25 |
26 | if __name__ == '__main__':
27 | parser = argparse.ArgumentParser()
28 | parser.add_argument('dataset', help='The dataset name')
29 | parser.add_argument('run_name', help='The folder name used to save model, output and evaluation metrics. This can be set to any word')
30 | parser.add_argument('--archive', type=str, required=True, help='The archive name that the dataset belongs to. This can be set to forecast_csv, or forecast_csv_univar')
31 | parser.add_argument('--gpu', type=int, default=0, help='The gpu no. used for training and inference (defaults to 0)')
32 | parser.add_argument('--batch-size', type=int, default=8, help='The batch size (defaults to 8)')
33 | parser.add_argument('--lr', type=float, default=0.001, help='The learning rate (defaults to 0.001)')
34 | parser.add_argument('--repr-dims', type=int, default=320, help='The representation dimension (defaults to 320)')
35 | parser.add_argument('--max-train-length', type=int, default=3000, help='For sequence with a length greater than , it would be cropped into some sequences, each of which has a length less than (defaults to 3000)')
36 | parser.add_argument('--iters', type=int, default=None, help='The number of iterations')
37 | parser.add_argument('--epochs', type=int, default=None, help='The number of epochs')
38 | parser.add_argument('--save-every', type=int, default=None, help='Save the checkpoint every iterations/epochs')
39 | parser.add_argument('--seed', type=int, default=None, help='The random seed')
40 | parser.add_argument('--max-threads', type=int, default=None, help='The maximum allowed number of threads used by this process')
41 | parser.add_argument('--eval', action="store_true", help='Whether to perform evaluation after training')
42 |
43 | parser.add_argument('--kernels', type=int, nargs='+', default=[1, 2, 4, 8, 16, 32, 64, 128], help='The kernel sizes used in the mixture of AR expert layers')
44 | parser.add_argument('--alpha', type=float, default=0.0005, help='Weighting hyperparameter for loss function')
45 |
46 | args = parser.parse_args()
47 |
48 | print("Dataset:", args.dataset)
49 | print("Arguments:", str(args))
50 |
51 | device = init_dl_program(args.gpu, seed=args.seed, max_threads=args.max_threads)
52 |
53 | if args.archive == 'forecast_csv':
54 | task_type = 'forecasting'
55 | data, train_slice, valid_slice, test_slice, scaler, pred_lens, n_covariate_cols = datautils.load_forecast_csv(args.dataset)
56 | train_data = data[:, train_slice]
57 | elif args.archive == 'forecast_csv_univar':
58 | task_type = 'forecasting'
59 | data, train_slice, valid_slice, test_slice, scaler, pred_lens, n_covariate_cols = datautils.load_forecast_csv(args.dataset, univar=True)
60 | train_data = data[:, train_slice]
61 | elif args.archive == 'forecast_npy':
62 | task_type = 'forecasting'
63 | data, train_slice, valid_slice, test_slice, scaler, pred_lens, n_covariate_cols = datautils.load_forecast_npy(args.dataset)
64 | train_data = data[:, train_slice]
65 | elif args.archive == 'forecast_npy_univar':
66 | task_type = 'forecasting'
67 | data, train_slice, valid_slice, test_slice, scaler, pred_lens, n_covariate_cols = datautils.load_forecast_npy(args.dataset, univar=True)
68 | train_data = data[:, train_slice]
69 | else:
70 | raise ValueError(f"Archive type {args.archive} is not supported.")
71 |
72 | config = dict(
73 | batch_size=args.batch_size,
74 | lr=args.lr,
75 | output_dims=args.repr_dims,
76 | )
77 |
78 | if args.save_every is not None:
79 | unit = 'epoch' if args.epochs is not None else 'iter'
80 | config[f'after_{unit}_callback'] = save_checkpoint_callback(args.save_every, unit)
81 |
82 | run_dir = f"training/{args.dataset}/{name_with_datetime(args.run_name)}"
83 |
84 | os.makedirs(run_dir, exist_ok=True)
85 |
86 | t = time.time()
87 |
88 | model = CoST(
89 | input_dims=train_data.shape[-1],
90 | kernels=args.kernels,
91 | alpha=args.alpha,
92 | max_train_length=args.max_train_length,
93 | device=device,
94 | **config
95 | )
96 |
97 | loss_log = model.fit(
98 | train_data,
99 | n_epochs=args.epochs,
100 | n_iters=args.iters,
101 | verbose=True
102 | )
103 | model.save(f'{run_dir}/model.pkl')
104 |
105 | t = time.time() - t
106 | print(f"\nTraining time: {datetime.timedelta(seconds=t)}\n")
107 |
108 | if args.eval:
109 | out, eval_res = tasks.eval_forecasting(model, data, train_slice, valid_slice, test_slice, scaler, pred_lens, n_covariate_cols, args.max_train_length-1)
110 | print('Evaluation result:', eval_res)
111 | pkl_save(f'{run_dir}/eval_res.pkl', eval_res)
112 | pkl_save(f'{run_dir}/out.pkl', out)
113 |
114 | print("Finished.")
115 |
--------------------------------------------------------------------------------
/models/encoder.py:
--------------------------------------------------------------------------------
1 | import math
2 | from typing import List
3 |
4 | import torch
5 | from torch import nn
6 | import torch.nn.functional as F
7 | import torch.fft as fft
8 | from einops import reduce, rearrange, repeat
9 |
10 | import numpy as np
11 |
12 | from .dilated_conv import DilatedConvEncoder
13 |
14 |
15 | def generate_continuous_mask(B, T, n=5, l=0.1):
16 | res = torch.full((B, T), True, dtype=torch.bool)
17 | if isinstance(n, float):
18 | n = int(n * T)
19 | n = max(min(n, T // 2), 1)
20 |
21 | if isinstance(l, float):
22 | l = int(l * T)
23 | l = max(l, 1)
24 |
25 | for i in range(B):
26 | for _ in range(n):
27 | t = np.random.randint(T-l+1)
28 | res[i, t:t+l] = False
29 | return res
30 |
31 |
32 | def generate_binomial_mask(B, T, p=0.5):
33 | return torch.from_numpy(np.random.binomial(1, p, size=(B, T))).to(torch.bool)
34 |
35 |
36 | class BandedFourierLayer(nn.Module):
37 | def __init__(self, in_channels, out_channels, band, num_bands, length=201):
38 | super().__init__()
39 |
40 | self.length = length
41 | self.total_freqs = (self.length // 2) + 1
42 |
43 | self.in_channels = in_channels
44 | self.out_channels = out_channels
45 |
46 | self.band = band # zero indexed
47 | self.num_bands = num_bands
48 |
49 | self.num_freqs = self.total_freqs // self.num_bands + (self.total_freqs % self.num_bands if self.band == self.num_bands - 1 else 0)
50 |
51 | self.start = self.band * (self.total_freqs // self.num_bands)
52 | self.end = self.start + self.num_freqs
53 |
54 |
55 | # case: from other frequencies
56 | self.weight = nn.Parameter(torch.empty((self.num_freqs, in_channels, out_channels), dtype=torch.cfloat))
57 | self.bias = nn.Parameter(torch.empty((self.num_freqs, out_channels), dtype=torch.cfloat))
58 | self.reset_parameters()
59 |
60 | def forward(self, input):
61 | # input - b t d
62 | b, t, _ = input.shape
63 | input_fft = fft.rfft(input, dim=1)
64 | output_fft = torch.zeros(b, t // 2 + 1, self.out_channels, device=input.device, dtype=torch.cfloat)
65 | output_fft[:, self.start:self.end] = self._forward(input_fft)
66 | return fft.irfft(output_fft, n=input.size(1), dim=1)
67 |
68 | def _forward(self, input):
69 | output = torch.einsum('bti,tio->bto', input[:, self.start:self.end], self.weight)
70 | return output + self.bias
71 |
72 | def reset_parameters(self) -> None:
73 | nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
74 | fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)
75 | bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
76 | nn.init.uniform_(self.bias, -bound, bound)
77 |
78 |
79 | class CoSTEncoder(nn.Module):
80 | def __init__(self, input_dims, output_dims,
81 | kernels: List[int],
82 | length: int,
83 | hidden_dims=64, depth=10,
84 | mask_mode='binomial'):
85 | super().__init__()
86 |
87 | component_dims = output_dims // 2
88 |
89 | self.input_dims = input_dims
90 | self.output_dims = output_dims
91 | self.component_dims = component_dims
92 | self.hidden_dims = hidden_dims
93 | self.mask_mode = mask_mode
94 | self.input_fc = nn.Linear(input_dims, hidden_dims)
95 |
96 | self.feature_extractor = DilatedConvEncoder(
97 | hidden_dims,
98 | [hidden_dims] * depth + [output_dims],
99 | kernel_size=3
100 | )
101 |
102 | self.repr_dropout = nn.Dropout(p=0.1)
103 |
104 | self.kernels = kernels
105 |
106 | self.tfd = nn.ModuleList(
107 | [nn.Conv1d(output_dims, component_dims, k, padding=k-1) for k in kernels]
108 | )
109 |
110 | self.sfd = nn.ModuleList(
111 | [BandedFourierLayer(output_dims, component_dims, b, 1, length=length) for b in range(1)]
112 | )
113 |
114 | def forward(self, x, tcn_output=False, mask='all_true'): # x: B x T x input_dims
115 | nan_mask = ~x.isnan().any(axis=-1)
116 | x[~nan_mask] = 0
117 | x = self.input_fc(x) # B x T x Ch
118 |
119 | # generate & apply mask
120 | if mask is None:
121 | if self.training:
122 | mask = self.mask_mode
123 | else:
124 | mask = 'all_true'
125 |
126 | if mask == 'binomial':
127 | mask = generate_binomial_mask(x.size(0), x.size(1)).to(x.device)
128 | elif mask == 'continuous':
129 | mask = generate_continuous_mask(x.size(0), x.size(1)).to(x.device)
130 | elif mask == 'all_true':
131 | mask = x.new_full((x.size(0), x.size(1)), True, dtype=torch.bool)
132 | elif mask == 'all_false':
133 | mask = x.new_full((x.size(0), x.size(1)), False, dtype=torch.bool)
134 | elif mask == 'mask_last':
135 | mask = x.new_full((x.size(0), x.size(1)), True, dtype=torch.bool)
136 | mask[:, -1] = False
137 |
138 | mask &= nan_mask
139 | x[~mask] = 0
140 |
141 | # conv encoder
142 | x = x.transpose(1, 2) # B x Ch x T
143 | x = self.feature_extractor(x) # B x Co x T
144 |
145 | if tcn_output:
146 | return x.transpose(1, 2)
147 |
148 | trend = []
149 | for idx, mod in enumerate(self.tfd):
150 | out = mod(x) # b d t
151 | if self.kernels[idx] != 1:
152 | out = out[..., :-(self.kernels[idx] - 1)]
153 | trend.append(out.transpose(1, 2)) # b t d
154 | trend = reduce(
155 | rearrange(trend, 'list b t d -> list b t d'),
156 | 'list b t d -> b t d', 'mean'
157 | )
158 |
159 | x = x.transpose(1, 2) # B x T x Co
160 |
161 | season = []
162 | for mod in self.sfd:
163 | out = mod(x) # b t d
164 | season.append(out)
165 | season = season[0]
166 |
167 | return trend, self.repr_dropout(season)
168 |
--------------------------------------------------------------------------------
/cost.py:
--------------------------------------------------------------------------------
1 | import sys, math, random, copy
2 | from typing import Union, Callable, Optional, List
3 |
4 | import torch
5 | import torch.nn as nn
6 | import torch.nn.functional as F
7 | import torch.fft as fft
8 | from torch.utils.data import TensorDataset, DataLoader, Dataset
9 |
10 | import numpy as np
11 | from einops import rearrange, repeat, reduce
12 |
13 | from models.encoder import CoSTEncoder
14 | from utils import take_per_row, split_with_nan, centerize_vary_length_series, torch_pad_nan
15 |
16 |
17 | class PretrainDataset(Dataset):
18 |
19 | def __init__(self,
20 | data,
21 | sigma,
22 | p=0.5,
23 | multiplier=10):
24 | super().__init__()
25 | self.data = data
26 | self.p = p
27 | self.sigma = sigma
28 | self.multiplier = multiplier
29 | self.N, self.T, self.D = data.shape # num_ts, time, dim
30 |
31 | def __getitem__(self, item):
32 | ts = self.data[item % self.N]
33 | return self.transform(ts), self.transform(ts)
34 |
35 | def __len__(self):
36 | return self.data.size(0) * self.multiplier
37 |
38 | def transform(self, x):
39 | return self.jitter(self.shift(self.scale(x)))
40 |
41 | def jitter(self, x):
42 | if random.random() > self.p:
43 | return x
44 | return x + (torch.randn(x.shape) * self.sigma)
45 |
46 | def scale(self, x):
47 | if random.random() > self.p:
48 | return x
49 | return x * (torch.randn(x.size(-1)) * self.sigma + 1)
50 |
51 | def shift(self, x):
52 | if random.random() > self.p:
53 | return x
54 | return x + (torch.randn(x.size(-1)) * self.sigma)
55 |
56 |
57 | class CoSTModel(nn.Module):
58 | def __init__(self,
59 | encoder_q: nn.Module, encoder_k: nn.Module,
60 | kernels: List[int],
61 | device: Optional[str] = 'cuda',
62 | dim: Optional[int] = 128,
63 | alpha: Optional[float] = 0.05,
64 | K: Optional[int] = 65536,
65 | m: Optional[float] = 0.999,
66 | T: Optional[float] = 0.07):
67 | super().__init__()
68 |
69 | self.K = K
70 | self.m = m
71 | self.T = T
72 | self.device = device
73 |
74 | self.kernels = kernels
75 |
76 | self.alpha = alpha
77 |
78 | self.encoder_q = encoder_q
79 | self.encoder_k = encoder_k
80 |
81 | # create the encoders
82 | self.head_q = nn.Sequential(
83 | nn.Linear(dim, dim),
84 | nn.ReLU(),
85 | nn.Linear(dim, dim)
86 | )
87 | self.head_k = nn.Sequential(
88 | nn.Linear(dim, dim),
89 | nn.ReLU(),
90 | nn.Linear(dim, dim)
91 | )
92 |
93 | for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
94 | param_k.data.copy_(param_q.data) # initialize
95 | param_k.requires_grad = False # not update by gradient
96 | for param_q, param_k in zip(self.head_q.parameters(), self.head_k.parameters()):
97 | param_k.data.copy_(param_q.data) # initialize
98 | param_k.requires_grad = False # not update by gradient
99 |
100 | self.register_buffer('queue', F.normalize(torch.randn(dim, K), dim=0))
101 | self.register_buffer('queue_ptr', torch.zeros(1, dtype=torch.long))
102 |
103 |
104 | def compute_loss(self, q, k, k_negs):
105 | # compute logits
106 | # positive logits: Nx1
107 | l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)
108 | # negative logits: NxK
109 | l_neg = torch.einsum('nc,ck->nk', [q, k_negs])
110 |
111 | # logits: Nx(1+K)
112 | logits = torch.cat([l_pos, l_neg], dim=1)
113 |
114 | # apply temperature
115 | logits /= self.T
116 |
117 | # labels: positive key indicators - first dim of each batch
118 | labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()
119 | loss = F.cross_entropy(logits, labels)
120 |
121 | return loss
122 |
123 | def convert_coeff(self, x, eps=1e-6):
124 | amp = torch.sqrt((x.real + eps).pow(2) + (x.imag + eps).pow(2))
125 | phase = torch.atan2(x.imag, x.real + eps)
126 | return amp, phase
127 |
128 | def instance_contrastive_loss(self, z1, z2):
129 | B, T = z1.size(0), z1.size(1)
130 | z = torch.cat([z1, z2], dim=0) # 2B x T x C
131 | z = z.transpose(0, 1) # T x 2B x C
132 | sim = torch.matmul(z, z.transpose(1, 2)) # T x 2B x 2B
133 | logits = torch.tril(sim, diagonal=-1)[:, :, :-1] # T x 2B x (2B-1)
134 | logits += torch.triu(sim, diagonal=1)[:, :, 1:]
135 | logits = -F.log_softmax(logits, dim=-1)
136 |
137 | i = torch.arange(B, device=z1.device)
138 | loss = (logits[:, i, B + i - 1].mean() + logits[:, B + i, i].mean()) / 2
139 | return loss
140 |
141 | def forward(self, x_q, x_k):
142 | # compute query features
143 | rand_idx = np.random.randint(0, x_q.shape[1])
144 |
145 | q_t, q_s = self.encoder_q(x_q)
146 | if q_t is not None:
147 | q_t = F.normalize(self.head_q(q_t[:, rand_idx]), dim=-1)
148 |
149 | # compute key features
150 | with torch.no_grad(): # no gradient for keys
151 | self._momentum_update_key_encoder() # update key encoder
152 | k_t, k_s = self.encoder_k(x_k)
153 | if k_t is not None:
154 | k_t = F.normalize(self.head_k(k_t[:, rand_idx]), dim=-1)
155 |
156 | loss = 0
157 |
158 | loss += self.compute_loss(q_t, k_t, self.queue.clone().detach())
159 | self._dequeue_and_enqueue(k_t)
160 |
161 | q_s = F.normalize(q_s, dim=-1)
162 | _, k_s = self.encoder_q(x_k)
163 | k_s = F.normalize(k_s, dim=-1)
164 |
165 | q_s_freq = fft.rfft(q_s, dim=1)
166 | k_s_freq = fft.rfft(k_s, dim=1)
167 | q_s_amp, q_s_phase = self.convert_coeff(q_s_freq)
168 | k_s_amp, k_s_phase = self.convert_coeff(k_s_freq)
169 |
170 | seasonal_loss = self.instance_contrastive_loss(q_s_amp, k_s_amp) + \
171 | self.instance_contrastive_loss(q_s_phase,k_s_phase)
172 | loss += (self.alpha * (seasonal_loss/2))
173 |
174 | return loss
175 |
176 | @torch.no_grad()
177 | def _momentum_update_key_encoder(self):
178 | """
179 | Momentum update for key encoder
180 | """
181 | for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):
182 | param_k.data = param_k.data * self.m + param_q.data * (1 - self.m)
183 | for param_q, param_k in zip(self.head_q.parameters(), self.head_k.parameters()):
184 | param_k.data = param_k.data * self.m + param_q.data * (1 - self.m)
185 |
186 | @torch.no_grad()
187 | def _dequeue_and_enqueue(self, keys):
188 | batch_size = keys.shape[0]
189 |
190 | ptr = int(self.queue_ptr)
191 | assert self.K % batch_size == 0
192 |
193 | # replace keys at ptr (dequeue and enqueue)
194 | self.queue[:, ptr:ptr + batch_size] = keys.T
195 |
196 | ptr = (ptr + batch_size) % self.K
197 | self.queue_ptr[0] = ptr
198 |
199 |
200 | class CoST:
201 | def __init__(self,
202 | input_dims: int,
203 | kernels: List[int],
204 | alpha: bool,
205 | max_train_length: int,
206 | output_dims: int = 320,
207 | hidden_dims: int = 64,
208 | depth: int = 10,
209 | device: 'str' ='cuda',
210 | lr: float = 0.001,
211 | batch_size: int = 16,
212 | after_iter_callback: Union[Callable, None] = None,
213 | after_epoch_callback: Union[Callable, None] = None):
214 |
215 | super().__init__()
216 | self.input_dims = input_dims
217 | self.output_dims = output_dims
218 | self.hidden_dims = hidden_dims
219 | self.device = device
220 | self.lr = lr
221 | self.batch_size = batch_size
222 | self.max_train_length = max_train_length
223 |
224 | if kernels is None:
225 | kernels = []
226 |
227 | self.net = CoSTEncoder(
228 | input_dims=input_dims, output_dims=output_dims,
229 | kernels=kernels,
230 | length=max_train_length,
231 | hidden_dims=hidden_dims, depth=depth,
232 | ).to(self.device)
233 |
234 | self.cost = CoSTModel(
235 | self.net,
236 | copy.deepcopy(self.net),
237 | kernels=kernels,
238 | dim=self.net.component_dims,
239 | alpha=alpha,
240 | K=256,
241 | device=self.device,
242 | ).to(self.device)
243 |
244 | self.after_iter_callback = after_iter_callback
245 | self.after_epoch_callback = after_epoch_callback
246 |
247 | self.n_epochs = 0
248 | self.n_iters = 0
249 |
250 | def fit(self, train_data, n_epochs=None, n_iters=None, verbose=False):
251 | assert train_data.ndim == 3
252 |
253 | if n_iters is None and n_epochs is None:
254 | n_iters = 200 if train_data.size <= 100000 else 600
255 |
256 | if self.max_train_length is not None:
257 | sections = train_data.shape[1] // self.max_train_length
258 | if sections >= 2:
259 | train_data = np.concatenate(split_with_nan(train_data, sections, axis=1), axis=0)
260 |
261 | temporal_missing = np.isnan(train_data).all(axis=-1).any(axis=0)
262 | if temporal_missing[0] or temporal_missing[-1]:
263 | train_data = centerize_vary_length_series(train_data)
264 |
265 | train_data = train_data[~np.isnan(train_data).all(axis=2).all(axis=1)]
266 |
267 | multiplier = 1 if train_data.shape[0] >= self.batch_size else math.ceil(self.batch_size / train_data.shape[0])
268 | train_dataset = PretrainDataset(torch.from_numpy(train_data).to(torch.float), sigma=0.5, multiplier=multiplier)
269 | train_loader = DataLoader(train_dataset, batch_size=min(self.batch_size, len(train_dataset)), shuffle=True, drop_last=True)
270 |
271 | optimizer = torch.optim.SGD([p for p in self.cost.parameters() if p.requires_grad],
272 | lr=self.lr,
273 | momentum=0.9,
274 | weight_decay=1e-4)
275 |
276 | loss_log = []
277 |
278 | while True:
279 | if n_epochs is not None and self.n_epochs >= n_epochs:
280 | break
281 |
282 | cum_loss = 0
283 | n_epoch_iters = 0
284 |
285 | interrupted = False
286 | for batch in train_loader:
287 | if n_iters is not None and self.n_iters >= n_iters:
288 | interrupted = True
289 | break
290 |
291 | x_q, x_k = map(lambda x: x.to(self.device), batch)
292 | if self.max_train_length is not None and x_q.size(1) > self.max_train_length:
293 | window_offset = np.random.randint(x_q.size(1) - self.max_train_length + 1)
294 | x_q = x_q[:, window_offset : window_offset + self.max_train_length]
295 | x_k = x_k[:, window_offset : window_offset + self.max_train_length]
296 |
297 | optimizer.zero_grad()
298 |
299 | loss = self.cost(x_q, x_k)
300 |
301 | loss.backward()
302 | optimizer.step()
303 |
304 | cum_loss += loss.item()
305 | n_epoch_iters += 1
306 |
307 | self.n_iters += 1
308 |
309 | if self.after_iter_callback is not None:
310 | self.after_iter_callback(self, loss.item())
311 |
312 | if n_iters is not None:
313 | adjust_learning_rate(optimizer, self.lr, self.n_iters, n_iters)
314 |
315 | if interrupted:
316 | break
317 |
318 | cum_loss /= n_epoch_iters
319 | loss_log.append(cum_loss)
320 | if verbose:
321 | print(f"Epoch #{self.n_epochs}: loss={cum_loss}")
322 | self.n_epochs += 1
323 |
324 | if self.after_epoch_callback is not None:
325 | self.after_epoch_callback(self, cum_loss)
326 |
327 | if n_epochs is not None:
328 | adjust_learning_rate(optimizer, self.lr, self.n_epochs, n_epochs)
329 |
330 | return loss_log
331 |
332 | def _eval_with_pooling(self, x, mask=None, slicing=None, encoding_window=None):
333 | out_t, out_s = self.net(x.to(self.device, non_blocking=True)) # l b t d
334 | out = torch.cat([out_t[:, -1], out_s[:, -1]], dim=-1)
335 | return rearrange(out.cpu(), 'b d -> b () d')
336 |
337 | def encode(self, data, mode, mask=None, encoding_window=None, casual=False, sliding_length=None, sliding_padding=0, batch_size=None):
338 | if mode == 'forecasting':
339 | encoding_window = None
340 | slicing = None
341 | else:
342 | raise NotImplementedError(f"mode {mode} has not been implemented")
343 |
344 | assert data.ndim == 3
345 | if batch_size is None:
346 | batch_size = self.batch_size
347 | n_samples, ts_l, _ = data.shape
348 |
349 | org_training = self.net.training
350 | self.net.eval()
351 |
352 | dataset = TensorDataset(torch.from_numpy(data).to(torch.float))
353 | loader = DataLoader(dataset, batch_size=batch_size)
354 |
355 | with torch.no_grad():
356 | output = []
357 | for batch in loader:
358 | x = batch[0]
359 | if sliding_length is not None:
360 | reprs = []
361 | if n_samples < batch_size:
362 | calc_buffer = []
363 | calc_buffer_l = 0
364 | for i in range(0, ts_l, sliding_length):
365 | l = i - sliding_padding
366 | r = i + sliding_length + (sliding_padding if not casual else 0)
367 | x_sliding = torch_pad_nan(
368 | x[:, max(l, 0) : min(r, ts_l)],
369 | left=-l if l<0 else 0,
370 | right=r-ts_l if r>ts_l else 0,
371 | dim=1
372 | )
373 | if n_samples < batch_size:
374 | if calc_buffer_l + n_samples > batch_size:
375 | out = self._eval_with_pooling(
376 | torch.cat(calc_buffer, dim=0),
377 | mask,
378 | slicing=slicing,
379 | encoding_window=encoding_window
380 | )
381 | reprs += torch.split(out, n_samples)
382 | calc_buffer = []
383 | calc_buffer_l = 0
384 | calc_buffer.append(x_sliding)
385 | calc_buffer_l += n_samples
386 | else:
387 | out = self._eval_with_pooling(
388 | x_sliding,
389 | mask,
390 | slicing=slicing,
391 | encoding_window=encoding_window
392 | )
393 | reprs.append(out)
394 |
395 | if n_samples < batch_size:
396 | if calc_buffer_l > 0:
397 | out = self._eval_with_pooling(
398 | torch.cat(calc_buffer, dim=0),
399 | mask,
400 | slicing=slicing,
401 | encoding_window=encoding_window
402 | )
403 | reprs += torch.split(out, n_samples)
404 | calc_buffer = []
405 | calc_buffer_l = 0
406 |
407 | out = torch.cat(reprs, dim=1)
408 | if encoding_window == 'full_series':
409 | out = F.max_pool1d(
410 | out.transpose(1, 2).contiguous(),
411 | kernel_size = out.size(1),
412 | ).squeeze(1)
413 | else:
414 | out = self._eval_with_pooling(x, mask, encoding_window=encoding_window)
415 | if encoding_window == 'full_series':
416 | out = out.squeeze(1)
417 |
418 | output.append(out)
419 |
420 | output = torch.cat(output, dim=0)
421 |
422 | self.net.train(org_training)
423 | return output.numpy()
424 |
425 | def save(self, fn):
426 | ''' Save the model to a file.
427 |
428 | Args:
429 | fn (str): filename.
430 | '''
431 | torch.save(self.net.state_dict(), fn)
432 |
433 | def load(self, fn):
434 | ''' Load the model from a file.
435 |
436 | Args:
437 | fn (str): filename.
438 | '''
439 | state_dict = torch.load(fn, map_location=self.device)
440 | self.net.load_state_dict(state_dict)
441 |
442 |
443 | def adjust_learning_rate(optimizer, lr, epoch, epochs):
444 | """Decay the learning rate based on schedule"""
445 | lr *= 0.5 * (1. + math.cos(math.pi * epoch / epochs))
446 | for param_group in optimizer.param_groups:
447 | param_group['lr'] = lr
448 |
--------------------------------------------------------------------------------
5 |