-
Notifications
You must be signed in to change notification settings - Fork 0
/
multihead_dual_tasks.py
108 lines (91 loc) · 3.99 KB
/
multihead_dual_tasks.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
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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
import argparse
from math import exp
import random
import numpy as np
import torch
import pandas as pd
from utils.config import fetch_configs
from datahandlers.cifar import SplitCIFARHandler
from net.smallconv import MultiHeadNet
from net.wideresnet import WideResNetMultihead
from utils.run_net import train, evaluate
SEED = 1234
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.backends.cudnn.deterministic = True
def run_experiment(exp_conf, gpu):
n = exp_conf['n']
hp = exp_conf['hp']
task_dict = exp_conf['task_dict']
df = pd.DataFrame()
for task in exp_conf['out_task']:
print("Doing task...T{}".format(task))
tasks = [task_dict[exp_conf['in_task']], task_dict[task]]
print("Tasks: {}".format(tasks))
dataset = SplitCIFARHandler(tasks)
i = 0
for mn in exp_conf['m_n_ratio']:
m = mn * n
print("m = {}".format(m))
for r, rep in enumerate(range(exp_conf['reps'])):
print("T{} vs. T{} : Doing rep...{}".format(exp_conf['in_task'], task, rep))
df.at[i, "m"] = mn
df.at[i, "r"] = r
dataset.sample_data(n=n, m=m, randomly=exp_conf['sample_scheme'])
train_loader = dataset.get_data_loader(hp['batch_size'], train=True)
if 'net' in exp_conf and exp_conf['net'] == 'wrn':
net = WideResNetMultihead(
depth=16,
num_task=len(tasks),
num_cls=len(tasks[0]),
widen_factor=4,
drop_rate=0.2
)
else:
net = MultiHeadNet(
num_task=len(tasks),
num_cls=len(tasks[0]),
channels=3,
avg_pool=2,
lin_size=320
)
optimizer = torch.optim.SGD(net.parameters(), lr=hp['lr'],
momentum=0.9, nesterov=True,
weight_decay=hp['l2_reg'])
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, hp['epochs'] * len(train_loader))
net = train(net, hp, train_loader, optimizer, lr_scheduler, gpu, verbose=False, task_id_flag=True)
risk = evaluate(net, dataset, gpu, task_id_flag=True)
print("Risk = %0.4f" % risk)
df.at[i, str(task)] = risk
i+=1
print("Saving individual results...")
df.to_csv('{}/{}_{}_{}_T{}_T{}.csv'.format(exp_conf['save_folder'], exp_conf['dataset'], exp_conf['net'], exp_conf['exp_name'], exp_conf['in_task'], task))
print("Saving bulk results...")
df.to_csv('{}/{}_{}_{}_T{}.csv'.format(exp_conf['save_folder'], exp_conf['dataset'], exp_conf['net'], exp_conf['exp_name'], exp_conf['in_task']))
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--exp_config', type=str,
default="./experiments/config/multihead_dual_tasks.yaml",
help="Experiment configuration")
parser.add_argument('--in_task', type=int,
help="Source task")
parser.add_argument('--out_task', nargs='+', type=int,
help="Target task(s)")
parser.add_argument('--gpu', type=str,
default='cuda:0',
help="GPU")
args = parser.parse_args()
exp_conf = fetch_configs(args.exp_config)
if args.in_task is not None:
exp_conf['in_task'] = args.in_task
if args.out_task is not None:
exp_conf['out_task'] = args.out_task
gpu = args.gpu
print("OOD Task(s) : {}".format(exp_conf['out_task']))
print("Target Task : {}".format(exp_conf['in_task']))
print("GPU : {}".format(gpu))
run_experiment(exp_conf, gpu)
if __name__ == '__main__':
main()