-
Notifications
You must be signed in to change notification settings - Fork 0
/
init_memory_bank.py
48 lines (37 loc) · 1.38 KB
/
init_memory_bank.py
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
import tensorflow as tf
from data_aug.transforms import rotate_transform, normalize
from memory_bank.memory_bank import MemoryBankTf
from models.baseline import CNN
tf.random.set_seed(99)
print(tf.test.is_gpu_available())
from keras.datasets import cifar10
# The data_aug, split between train and test sets:
(x_train, y_train), (x_test, y_test) = cifar10.load_data()
print('x_train shape:', x_train.shape)
print(x_train.shape[0], 'train samples')
print(x_test.shape[0], 'test samples')
indices = list(range(len(x_train)))
INPUT_SHAPE = (32, 32, 3)
EPOCHS = 1
N_NEGATIVES = 128
encoder = CNN(INPUT_SHAPE)
# _ = encoder(np.random.rand(1,32,32,3).astype(np.float32))
# encoder.load_weights('encoder.h5')
# create a dataset to initialize the Memory bank
memory_bank = MemoryBankTf(shape=(x_train.shape[0], 128), from_pickle=False)
dataset = tf.data.Dataset.from_tensor_slices((indices, x_train))
dataset = dataset.map(rotate_transform())
dataset = dataset.map(normalize)
dataset = dataset.repeat(EPOCHS)
dataset = dataset.batch(256)
# fill up the memory bank
counter = 0
for curr_indices, I, It in dataset:
# print(np.min(I), np.max(I))
# print(np.min(It), np.max(It))
_, f_vi = encoder(I, head=tf.constant('f'))
memory_bank.init_memory_bank(curr_indices, batch_features=f_vi)
counter += 1
print("Total iter:", counter)
memory_bank.save_memory_bank()
print("Memory bank is filled!")