-
Notifications
You must be signed in to change notification settings - Fork 32
/
train_gra_transf_inpt5_new_dropout_2layerMLP_fully_connected_graph_early_stop.py
371 lines (252 loc) · 14.2 KB
/
train_gra_transf_inpt5_new_dropout_2layerMLP_fully_connected_graph_early_stop.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
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
import argparse
import collections
import datetime
import os
import pickle
import time
import ipdb
import matplotlib.pyplot as plt
import numpy as np
import random
import torch
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from tensorboardX import SummaryWriter
from torch.autograd import Variable
from torch.nn.modules.module import Module
from torch.optim import lr_scheduler
from torch.utils.data import DataLoader
from dataloader.QuickdrawDataset4dict_fully_connected_graph_attention_mask import *
from network.gra_transf_inpt5_new_dropout_2layerMLP import *
from utils.AverageMeter import AverageMeter
from utils.Logger import Logger
from utils.accuracy import *
from utils.EarlyStopping import *
from tqdm import tqdm
################################################
# This python file contains four parts:
#
# Part 1. Argument Parser
# Part 2. configurations:
# Part 2-1. Basic configuration
# Part 2-2. dataloader instantiation
# Part 2-3. log configuration
# Part 2-4. configurations for loss function, network, and optimizer
# Part 3. 'train' function
# Part 4. 'validate' function
# Part 5. 'main' function
################################################
# Part 1. Argument Parser
parser = argparse.ArgumentParser(description='MGT_stage_1')
parser.add_argument("--exp", type=str, default="train_gra_transf_inpt5_new_dropout_2layerMLP_fully_connected_graph_early_stop_004", help="experiment")
parser.add_argument("--train_coordinate_path_root", type=str, default="/home/peng/dataset/tiny_quickdraw_coordinate/train/", help="train_sketch_coordinate_dir")
parser.add_argument("--val_coordinate_path_root", type=str, default="/home/peng/dataset/tiny_quickdraw_coordinate/val/", help="val_sketch_coordinate_dir")
parser.add_argument("--test_coordinate_path_root", type=str, default="/home/peng/dataset/tiny_quickdraw_coordinate/test/", help="test_sketch_coordinate_dir")
parser.add_argument("--sketch_list", type=str, default="./dataloader/tiny_train_set.txt", help="sketch_list_urls")
parser.add_argument("--sketch_list_4_val", type=str, default="./dataloader/tiny_val_set.txt", help="sketch_list_urls_4_validation")
parser.add_argument("--sketch_list_4_test", type=str, default="./dataloader/tiny_test_set.txt", help="sketch_list_urls_4_test")
parser.add_argument("--batch_size", type=int, default=512, help="batch_size")
parser.add_argument("--num_workers", type=int, default=12, help="num_workers")
parser.add_argument('--gpu', type=str, default="0", help='choose GPU')
args = parser.parse_args()
# Part 2. configurations
# Part 2-1. Basic configuration
basic_configs = collections.OrderedDict()
basic_configs['serial_number'] = args.exp
basic_configs['random_seed'] = int(time.time())
_seed = basic_configs['random_seed']
random.seed(_seed)
np.random.seed(_seed)
torch.manual_seed(_seed)
torch.cuda.manual_seed(_seed)
torch.cuda.manual_seed_all(_seed)
os.environ['PYTHONHASHSEED'] = str(_seed)
basic_configs['learning_rate'] = 0.00005
basic_configs['num_epochs'] = 100
basic_configs['early_stopping_patience'] = 10
# 001
#basic_configs["lr_protocol"] = [(5,0.00005), (10,0.00005 * 0.7), (15,0.00005 * (0.7 ** 2)), (20,0.00005 * (0.7 ** 3)),(25,0.00005 * (0.7 ** 4)), (30,0.00005 * (0.7 ** 5)), (35,0.00005 * (0.7 ** 6)), (40,0.00005 * (0.7 ** 7)), (45,0.00005 * (0.7 ** 8)), (50,0.00005 * (0.7 ** 9)), (55,0.00005 * (0.7 ** 10)), (60,0.00005 * (0.7 ** 11)), (65,0.00005 * (0.7 ** 12)), (70,0.00005 * (0.7 ** 13)), (75,0.00005 * (0.7 ** 14)), (80,0.00005 * (0.7 ** 15)), (85,0.00005 * (0.7 ** 16)), (90,0.00005 * (0.7 ** 17)), (95,0.00005 * (0.7 ** 18)), (100,0.00005 * (0.7 ** 19))]
# 004
basic_configs["lr_protocol"] = [(10,0.00005), (20,0.00005 * 0.7), (30,0.00005 * 0.7 * 0.7), (40,0.00005 * 0.7 * 0.7 * 0.7), (50,0.00005 * 0.7 * 0.7 * 0.7 * 0.7), (60,0.00005 * 0.7 * 0.7 * 0.7 * 0.7 * 0.7), (70,0.00005 * 0.7 * 0.7 * 0.7 * 0.7 * 0.7 * 0.7), (80,0.00005 * 0.7 * 0.7 * 0.7 * 0.7 * 0.7 * 0.7 * 0.7), (90,0.00005 * 0.7 * 0.7 * 0.7 * 0.7 * 0.7 * 0.7 * 0.7 * 0.7), (100,0.00005 * 0.7 * 0.7 * 0.7 * 0.7 * 0.7 * 0.7 * 0.7 * 0.7 * 0.7)]
basic_configs["display_step"] = 100
lr_protocol = basic_configs["lr_protocol"]
# Part 2-2. dataloader instantiation
dataloader_configs = collections.OrderedDict()
dataloader_configs['train_coordinate_path_root'] = args.train_coordinate_path_root
dataloader_configs['val_coordinate_path_root'] = args.val_coordinate_path_root
dataloader_configs['test_coordinate_path_root'] = args.test_coordinate_path_root
dataloader_configs['sketch_list'] = args.sketch_list
dataloader_configs['sketch_list_4_val'] = args.sketch_list_4_val
dataloader_configs['sketch_list_4_test'] = args.sketch_list_4_test
dataloader_configs['batch_size'] = args.batch_size
dataloader_configs['num_workers'] = args.num_workers
dataloader_configs['data_dict_4_train'] = './dataloader/tiny_train_dataset_dict.pickle'
data_dict_4_train_f = open(dataloader_configs['data_dict_4_train'], 'rb')
data_dict_4_train = pickle.load(data_dict_4_train_f)
dataloader_configs['data_dict_4_validation'] = './dataloader/tiny_val_dataset_dict.pickle'
data_dict_4_validation_f = open(dataloader_configs['data_dict_4_validation'], 'rb')
data_dict_4_validation = pickle.load(data_dict_4_validation_f)
dataloader_configs['data_dict_4_test'] = './dataloader/tiny_test_dataset_dict.pickle'
data_dict_4_test_f = open(dataloader_configs['data_dict_4_test'], 'rb')
data_dict_4_test = pickle.load(data_dict_4_test_f)
#ipdb.set_trace()
# create dataset
# -----------------------------------------------------------------------------------------------------
train_dataset = QuickdrawDataset_fully_connected_graph_attmask(dataloader_configs['train_coordinate_path_root'], dataloader_configs['sketch_list'], data_dict_4_train)
train_loader = DataLoader(train_dataset, batch_size=dataloader_configs['batch_size'], shuffle=True, num_workers=dataloader_configs['num_workers'])
val_dataset = QuickdrawDataset_fully_connected_graph_attmask(dataloader_configs['val_coordinate_path_root'], dataloader_configs['sketch_list_4_val'], data_dict_4_validation)
val_loader = DataLoader(val_dataset, batch_size=dataloader_configs['batch_size'], shuffle=False, num_workers=dataloader_configs['num_workers'])
test_dataset = QuickdrawDataset_fully_connected_graph_attmask(dataloader_configs['test_coordinate_path_root'], dataloader_configs['sketch_list_4_test'], data_dict_4_test)
test_loader = DataLoader(test_dataset, batch_size=dataloader_configs['batch_size'], shuffle=False, num_workers=dataloader_configs['num_workers'])
# Part 2-3. log configuration
exp_dir = os.path.join('./experimental_results', args.exp)
exp_log_dir = os.path.join(exp_dir, "log")
if not os.path.exists(exp_log_dir):
os.makedirs(exp_log_dir)
exp_visual_dir = os.path.join(exp_dir, "visual")
if not os.path.exists(exp_visual_dir):
os.makedirs(exp_visual_dir)
exp_ckpt_dir = os.path.join(exp_dir, "checkpoints")
if not os.path.exists(exp_ckpt_dir):
os.makedirs(exp_ckpt_dir)
now_str = datetime.datetime.now().__str__().replace(' ', '_')
writer_path = os.path.join(exp_visual_dir, now_str)
writer = SummaryWriter(writer_path)
logger_path = os.path.join(exp_log_dir, now_str + ".log")
logger = Logger(logger_path).get_logger()
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
logger.info("argument parser settings: {}".format(args))
logger.info("basic configuration settings: {}".format(basic_configs))
# Part 2-4. configurations for loss function, network, and optimizer
loss_function = nn.CrossEntropyLoss()
max_val_acc = 0.0
max_val_acc_epoch = -1
network_configs=collections.OrderedDict()
network_configs['output_dim']=345
network_configs['n_heads']=8
network_configs['embed_dim']=256
network_configs['n_layers']=4
network_configs['feed_forward_hidden']=4*network_configs['embed_dim']
network_configs['normalization']='batch'
network_configs['dropout']=0.10
network_configs['mlp_classifier_dropout']=0.10
logger.info("network configuration settings: {}".format(network_configs))
net = make_model(n_classes=345, coord_input_dim=2, feat_input_dim=2, feat_dict_size=103,
n_layers=network_configs['n_layers'], n_heads=network_configs['n_heads'],
embed_dim=network_configs['embed_dim'], feedforward_dim=network_configs['feed_forward_hidden'],
normalization=network_configs['normalization'], dropout=network_configs['dropout'], mlp_classifier_dropout=network_configs['mlp_classifier_dropout'])
net = net.cuda()
optimizer = torch.optim.Adam(net.parameters(), lr=0.00005)
# Part 3. 'train' function
def train_function(epoch):
training_loss = AverageMeter()
training_acc = AverageMeter()
net.train()
lr = next((lr for (max_epoch, lr) in lr_protocol if max_epoch > epoch), lr_protocol[-1][1])
for param_group in optimizer.param_groups:
param_group['lr'] = lr
logger.info("set learning rate to: {}".format(lr))
for idx, (coordinate, label, flag_bits, stroke_len, attention_mask, padding_mask, position_encoding) in enumerate(tqdm(train_loader, ascii=True)):
coordinate = coordinate.cuda()
label = label.cuda()
flag_bits = flag_bits.cuda()
stroke_len = stroke_len.cuda()
attention_mask = attention_mask.cuda()
padding_mask = padding_mask.cuda()
position_encoding = position_encoding.cuda()
# Resize inputs
flag_bits.squeeze_(2)
position_encoding.squeeze_(2)
stroke_len.unsqueeze_(1)
optimizer.zero_grad()
output = net(coordinate, flag_bits, position_encoding, attention_mask, padding_mask, stroke_len)
# ipdb.set_trace()
batch_loss = loss_function(output, label)
batch_loss.backward()
optimizer.step()
training_loss.update(batch_loss.item(), coordinate.size(0))
training_acc.update(accuracy(output, label, topk = (1,))[0].item(), coordinate.size(0))
if (idx + 1) % basic_configs["display_step"] == 0:
logger.info(
"==> Iteration [{}][{}/{}]:".format(epoch + 1, idx + 1, len(train_loader)))
logger.info("current batch loss: {}".format(
batch_loss.item()))
logger.info("average loss: {}".format(
training_loss.avg))
logger.info("average acc: {}".format(training_acc.avg))
logger.info("Begin evaluating on validation set")
validation_loss, validation_acc = validate_function(val_loader)
logger.info("Begin evaluating on testing set")
test_loss, test_acc = validate_function(test_loader)
writer.add_scalars("loss", {
"training_loss":training_loss.avg,
"validation_loss":validation_loss.avg,
"test_loss":test_loss.avg
}, epoch+1)
writer.add_scalars("acc", {
"training_acc":training_acc.avg,
"validation_acc":validation_acc.avg,
"test_acc":test_acc.avg,
}, epoch+1)
return validation_acc
# Part 4. 'validate' function
def validate_function(data_loader):
validation_loss = AverageMeter()
validation_acc_1 = AverageMeter()
validation_acc_5 = AverageMeter()
validation_acc_10 = AverageMeter()
net.eval()
timelist = list()
with torch.no_grad():
for idx, (coordinate, label, flag_bits, stroke_len, attention_mask, padding_mask, position_encoding) in enumerate(tqdm(data_loader, ascii=True)):
coordinate = coordinate.cuda()
label = label.cuda()
flag_bits = flag_bits.cuda()
stroke_len = stroke_len.cuda()
attention_mask = attention_mask.cuda()
padding_mask = padding_mask.cuda()
position_encoding = position_encoding.cuda()
# Resize inputs
flag_bits.squeeze_(2)
position_encoding.squeeze_(2)
stroke_len.unsqueeze_(1)
tic = time.time()
output = net(coordinate, flag_bits, position_encoding, attention_mask, padding_mask, stroke_len)
timelist.append(time.time() - tic)
batch_loss = loss_function(output, label)
validation_loss.update(batch_loss.item(), coordinate.size(0))
acc_1, acc_5, acc_10 = accuracy(output, label, topk = (1, 5, 10))
validation_acc_1.update(acc_1, coordinate.size(0))
validation_acc_5.update(acc_5, coordinate.size(0))
validation_acc_10.update(acc_10, coordinate.size(0))
logger.info("==> Evaluation Result: ")
logger.info("loss: {} acc@1: {} acc@5: {} acc@10: {}".format(validation_loss.avg, validation_acc_1.avg, validation_acc_5.avg, validation_acc_10.avg))
logger.info("Total inference time: {}s".format(sum(timelist)))
return validation_loss, validation_acc_1
# Part 5. 'main' function
if __name__ == '__main__':
logger.info("Begin evaluating on validation set before training")
validate_function(val_loader)
logger.info("training status: ")
early_stopping = EarlyStopping(patience=basic_configs['early_stopping_patience'], delta=0)
for epoch in range(basic_configs['num_epochs']):
logger.info("Begin training epoch {}".format(epoch + 1))
validation_acc = train_function(epoch)
if validation_acc.avg > max_val_acc:
max_val_acc = validation_acc.avg
max_val_acc_epoch = epoch + 1
early_stopping(validation_acc.avg)
logger.info("Early stopping counter: {}".format(early_stopping.counter))
logger.info("Early stopping best_score: {}".format(early_stopping.best_score))
logger.info("Early stopping early_stop: {}".format(early_stopping.early_stop))
if early_stopping.early_stop == True:
logger.info("Early stopping after Epoch: {}".format(epoch + 1))
break
net_checkpoint_name = args.exp + "_net_epoch" + str(epoch + 1)
net_checkpoint_path = os.path.join(exp_ckpt_dir, net_checkpoint_name)
net_state = {"epoch": epoch + 1,
"network": net.state_dict()}
torch.save(net_state, net_checkpoint_path)
logger.info("max_val_acc: {} max_val_acc_epoch: {}".format(max_val_acc, max_val_acc_epoch))