53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
import numpy as np
|
|
import pickle
|
|
import pycuda.driver as cuda
|
|
import sys
|
|
import os
|
|
from datetime import datetime
|
|
|
|
current_dir = os.path.dirname(os.path.abspath("__file__"))
|
|
parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
|
|
output_dir = os.path.join(parent_dir, "output")
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
sys.path.append(parent_dir)
|
|
|
|
cuda.init()
|
|
|
|
context = cuda.Device(3).make_context()
|
|
|
|
DATA_TYPE = np.float32
|
|
|
|
from env_manifold import CustomEnv
|
|
|
|
context.push()
|
|
env = CustomEnv(device_id=3)
|
|
context.pop()
|
|
|
|
size = [10, 10, 10]
|
|
|
|
def generate_random_group(size, low, high):
|
|
intervals = np.linspace(low, high, size + 1)
|
|
group = np.concatenate([np.random.uniform(intervals[i], intervals[i+1], 1) for i in range(size)])
|
|
return np.sort(group)
|
|
|
|
group1 = generate_random_group(size[0], -1, 1)
|
|
group2 = generate_random_group(size[1], -1, 1)
|
|
group3 = generate_random_group(size[2], -1, 1)
|
|
|
|
data = np.empty(size, dtype=object)
|
|
for i, a1 in enumerate(sorted(group1)):
|
|
for j, a2 in enumerate(sorted(group2)):
|
|
for k, a3 in enumerate(sorted(group3)):
|
|
context.push()
|
|
action = np.array([a1, a2, a3], dtype=np.float32)
|
|
env.reset()
|
|
for _ in range(400):
|
|
_, _, _, _, _ = env.step(action)
|
|
context.pop()
|
|
fifo = np.array(env.fifo_states.copy())
|
|
data[i, j, k] = {'action': action, 'fifo': fifo}
|
|
current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|
print(f"{current_time} - ({i}, {j}, {k})")
|
|
|
|
with open(os.path.join(output_dir, "manifold_1k.pkl"), 'wb') as f:
|
|
pickle.dump(data, f) |