-
Notifications
You must be signed in to change notification settings - Fork 1
/
orchestrate_results.py
56 lines (50 loc) · 2.06 KB
/
orchestrate_results.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
49
50
51
52
53
54
55
56
import multiprocessing
from multiprocessing import Process, Lock
from refactored_hyperparam_search import get_results
import time
import torch
NUM_GPUS = 8
def f(gpu_locks, args):
while True: # Infinite loop until it finds a GPU break out
for gpu_idx, gpu_lock in enumerate(gpu_locks):
acq = gpu_lock.acquire(block=False)
if acq:
import os
os.environ['CUDA_VISIBLE_DEVICES']=str(gpu_idx)
torch.cuda.set_device(gpu_idx)
print('Using gpu:', gpu_idx, 'with args', args)
get_results(**args)
print('Returning gpu:', gpu_idx)
gpu_lock.release()
return
time.sleep(100) # Wait 100 seconds before looping again to conserve CPU resources
"""
VLM_ARG, CLASSIFIER_ARG, dataset_name=["smsm", "kinetics_100"],
num_shots = [1, 2, 4, 8, 16], num_episodes = [4]
"""
if __name__ == '__main__':
multiprocessing.set_start_method('spawn')
gpu_locks = [Lock() for _ in range(NUM_GPUS)]
classifier_arg = "cona"
vlms = ["videoclip", "clip", "miles"] # Order slowest to fastest
num_episodes = [2, 4, 4] # Corresponds to ^
datasets = ["kinetics_100", "moma_act"] # Order largest to smallest
num_shots = [16, 8, 4, 2, 1] # Order largest to smallest
default_args_to_run = [] # Running with CoNa
for vlm, n_episodes in zip(vlms, num_episodes):
for dataset in datasets:
for n_shots in num_shots:
default_args_to_run.append({
"VLM_ARG":vlm,
"CLASSIFIER_ARG":classifier_arg,
"dataset_name":[dataset],
"num_shots":[n_shots],
"num_episodes":[n_episodes],
})
processes = []
for args in default_args_to_run:
processes.append(Process(target=f, args=(gpu_locks, args)))
processes[-1].start()
time.sleep(1) # Allow enough time for first processes to get GPUs first
for process in processes:
process.join()