InContextContinualLearning/slow_data_generation.py at main · Sequential-Intelligence-Lab/InContextContinualLearning · 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
import os
from argparse import ArgumentParser, Namespace

import torch
from tqdm import tqdm

from supervised.model import FixLTUNet
from utils import set_seed


def generate_problem_data(
        flip_after: int,
        num_data_points: int,
        num_inputs: int,
        seed: int,
        data_path: str = 'data/slow',
        num_target_features: int = 100,
        num_flipping_bits: int = 15,
        beta: float = 0.7,
        flip_one: bool = True):
    """
    Generates data for one run on the slowly changing regression problem
    """
    os.makedirs(data_path, exist_ok=True)

    set_seed(seed)

    target_network = FixLTUNet(
        num_inputs=num_inputs,
        num_features=num_target_features,
        beta=beta,
    )

    num_flips = int(num_data_points/flip_after)
    if num_data_points % flip_after != 0:
        num_flips += 1  # to ensure the number of flips is correct
    num_data_points = num_flips * flip_after
    flipping_bits = torch.randint(2, size=(num_flips, num_flipping_bits),
                                  dtype=torch.float32)
    if num_flipping_bits > 0:
        if flip_one:
            for i in range(1, num_flips):
                flipping_bits[i] = flipping_bits[i-1]
                bit_to_flip = torch.randint(num_flipping_bits, (1, ))
                flipping_bits[i][bit_to_flip] = 1 - \
                    flipping_bits[i-1][bit_to_flip]

        flipping_bits = flipping_bits.repeat_interleave(flip_after, dim=0)
        random_bits = torch.randint(2, size=(num_data_points, num_inputs - num_flipping_bits),
                                    dtype=torch.float32)
        X = torch.cat((flipping_bits, random_bits), dim=1)
    else:
        X = torch.randint(2, size=(num_data_points, num_inputs),
                          dtype=torch.float32)

    Y = torch.zeros((num_flips, flip_after, 1), dtype=torch.float)
    X = X.reshape((num_flips, flip_after, num_inputs))
    with torch.no_grad():
        for i in tqdm(range(X.shape[0])):
            Y[i], _ = target_network.predict(x=X[i])

    data = dict(X=X, Y=Y)
    torch.save(data, os.path.join(data_path, f'{seed}.pt'))


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--flip_after', type=int,
                        help='number of data points before flipping', default=10000)
    parser.add_argument('--num_data_points', type=int,
                        help='number of data points', default=3e7)
    parser.add_argument('-d', '--dim_feature', type=int,
                        help='feature dimension', default=20)
    parser.add_argument('--seed', type=int, nargs='+',
                        help='random seed', default=list(range(10)))
    args: Namespace = parser.parse_args()
    print('Generating slowly changing regression data...')
    for s in args.seed:
        generate_problem_data(flip_after=args.flip_after,
                              num_data_points=args.num_data_points,
                              num_inputs=args.dim_feature,
                              seed=s)