-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.py
153 lines (130 loc) · 4.47 KB
/
main.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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
import random
import numpy as np
import torch
import argparse
from src.utils import *
from src import train
parser = argparse.ArgumentParser(description="Missing modality")
# Tasks
parser.add_argument(
"--pretrained_model",
type=str,
default=None,
help="name of the model to use (Transformer, etc.)",
)
parser.add_argument(
"--dataset",
type=str,
default="mosi",
help="dataset to use (mosei, mosi, iemocap, sims)",
)
parser.add_argument(
"--data_path",
type=str,
default=None,
help="path for storing the dataset",
)
# Dropouts
parser.add_argument("--attn_dropout", type=float, default=0.1, help="attention dropout")
parser.add_argument(
"--attn_dropout_a", type=float, default=0.1, help="attention dropout (for audio)"
)
parser.add_argument(
"--attn_dropout_v", type=float, default=0.1, help="attention dropout (for visual)"
)
parser.add_argument("--relu_dropout", type=float, default=0.1, help="relu dropout")
parser.add_argument(
"--embed_dropout", type=float, default=0.25, help="embedding dropout"
)
parser.add_argument(
"--res_dropout", type=float, default=0.1, help="residual block dropout"
)
parser.add_argument(
"--out_dropout", type=float, default=0.1, help="output layer dropout"
)
# Architecture
parser.add_argument(
"--nlevels", type=int, default=5, help="number of layers in the network"
)
parser.add_argument(
"--proj_dim", type=int, default=30, help="projection dimension of the network"
)
parser.add_argument(
"--num_heads",
type=int,
default=5,
help="number of heads for the transformer network",
)
parser.add_argument(
"--attn_mask", action="store_false", help="use attention mask for Transformer"
)
parser.add_argument("--prompt_dim", type=int, default=30)
parser.add_argument("--prompt_length", type=int, default=16)
# Tuning
parser.add_argument(
"--batch_size", type=int, default=64, metavar="N", help="batch size"
)
parser.add_argument("--clip", type=float, default=0.8, help="gradient clip value")
parser.add_argument("--lr", type=float, default=1e-3, help="initial learning rate")
parser.add_argument("--optim", type=str, default="Adam", help="optimizer to use")
parser.add_argument("--num_epochs", type=int, default=40, help="number of epochs")
parser.add_argument("--when", type=int, default=10, help="when to decay learning rate")
parser.add_argument("--drop_rate", type=float, default=0.6)
# Logistics
parser.add_argument(
"--log_interval",
type=int,
default=30,
help="frequency of result logging (default: 30)",
)
parser.add_argument("--seed", type=int, default=666, help="random seed")
parser.add_argument("--no_cuda", action="store_true", help="do not use cuda")
parser.add_argument("--name", type=str, default=None, help="name of the trial")
args = parser.parse_args()
dataset = str.lower(args.dataset.strip())
output_dim_dict = {"mosi": 1, "mosei": 1, "sims": 1, "iemocap": 4}
criterion_dict = {"iemocap": "CrossEntropyLoss"}
def setup_seed(seed):
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
np.random.seed(seed)
random.seed(seed)
torch.backends.cudnn.deterministic = True
torch.set_default_tensor_type("torch.FloatTensor")
use_cuda = False
if torch.cuda.is_available():
if args.no_cuda:
print(
"WARNING: You have a CUDA device, so you should probably not run with --no_cuda"
)
else:
torch.cuda.manual_seed(args.seed)
torch.set_default_tensor_type("torch.cuda.FloatTensor")
use_cuda = True
setup_seed(args.seed)
####################################################################
#
# Load the dataset
#
####################################################################
dataloaders, orig_dims, n_nums, seq_len = get_loader(args)
trainloder = dataloaders["train"]
validloder = dataloaders["valid"]
testloder = dataloaders["test"]
####################################################################
#
# Hyperparameters
#
####################################################################
hyp_params = args
hyp_params.orig_d_l, hyp_params.orig_d_a, hyp_params.orig_d_v = orig_dims
hyp_params.layers = args.nlevels
hyp_params.use_cuda = use_cuda
hyp_params.dataset = dataset
hyp_params.when = args.when
hyp_params.n_train, hyp_params.n_valid, hyp_params.n_test = n_nums
hyp_params.output_dim = output_dim_dict.get(dataset, 1)
hyp_params.criterion = criterion_dict.get(dataset, "L1Loss")
hyp_params.seq_len = seq_len
if __name__ == "__main__":
train.initiate(hyp_params, trainloder, validloder, testloder)