Superxixixi commited on
Commit
b89c008
·
1 Parent(s): fe7cbab

Delete train.py

Browse files
Files changed (1) hide show
  1. train.py +0 -197
train.py DELETED
@@ -1,197 +0,0 @@
1
- import time, os, torch, argparse, warnings, glob, pandas, json
2
-
3
- from utils.tools import *
4
- from dlhammer import bootstrap
5
- os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
6
- import torch.multiprocessing as mp
7
- import torch.distributed as dist
8
-
9
- from xxlib.utils.distributed import all_gather, all_reduce
10
- from torch import nn
11
- from dataLoader_multiperson import train_loader, val_loader
12
-
13
- from loconet import loconet
14
-
15
-
16
- class MyCollator(object):
17
-
18
- def __init__(self, cfg):
19
- self.cfg = cfg
20
-
21
- def __call__(self, data):
22
- audiofeatures = [item[0] for item in data]
23
- visualfeatures = [item[1] for item in data]
24
- labels = [item[2] for item in data]
25
- masks = [item[3] for item in data]
26
- cut_limit = self.cfg.MODEL.CLIP_LENGTH
27
- # pad audio
28
- lengths = torch.tensor([t.shape[1] for t in audiofeatures])
29
- max_len = max(lengths)
30
- padded_audio = torch.stack([
31
- torch.cat([i, i.new_zeros((i.shape[0], max_len - i.shape[1], i.shape[2]))], 1)
32
- for i in audiofeatures
33
- ], 0)
34
-
35
- if max_len > cut_limit * 4:
36
- padded_audio = padded_audio[:, :, :cut_limit * 4, ...]
37
-
38
- # pad video
39
- lengths = torch.tensor([t.shape[1] for t in visualfeatures])
40
- max_len = max(lengths)
41
- padded_video = torch.stack([
42
- torch.cat(
43
- [i, i.new_zeros((i.shape[0], max_len - i.shape[1], i.shape[2], i.shape[3]))], 1)
44
- for i in visualfeatures
45
- ], 0)
46
- padded_labels = torch.stack(
47
- [torch.cat([i, i.new_zeros((i.shape[0], max_len - i.shape[1]))], 1) for i in labels], 0)
48
- padded_masks = torch.stack(
49
- [torch.cat([i, i.new_zeros((i.shape[0], max_len - i.shape[1]))], 1) for i in masks], 0)
50
-
51
- if max_len > cut_limit:
52
- padded_video = padded_video[:, :, :cut_limit, ...]
53
- padded_labels = padded_labels[:, :, :cut_limit, ...]
54
- padded_masks = padded_masks[:, :, :cut_limit, ...]
55
-
56
- return padded_audio, padded_video, padded_labels, padded_masks
57
-
58
-
59
- class DataPrep():
60
-
61
- def __init__(self, cfg, world_size, rank):
62
- self.cfg = cfg
63
- self.world_size = world_size
64
- self.rank = rank
65
-
66
- def train_dataloader(self):
67
-
68
- loader = train_loader(self.cfg, trialFileName = self.cfg.trainTrialAVA, \
69
- audioPath = os.path.join(self.cfg.audioPathAVA , 'train'), \
70
- visualPath = os.path.join(self.cfg.visualPathAVA, 'train'), \
71
- num_speakers=self.cfg.MODEL.NUM_SPEAKERS,
72
- )
73
- train_sampler = torch.utils.data.distributed.DistributedSampler(
74
- loader, num_replicas=self.world_size, rank=self.rank)
75
- collator = MyCollator(self.cfg)
76
- trainLoader = torch.utils.data.DataLoader(loader,
77
- batch_size=self.cfg.TRAIN.BATCH_SIZE,
78
- pin_memory=False,
79
- num_workers=self.cfg.NUM_WORKERS,
80
- collate_fn=collator,
81
- sampler=train_sampler)
82
- return trainLoader
83
-
84
- def val_dataloader(self):
85
- loader = val_loader(self.cfg, trialFileName = self.cfg.evalTrialAVA, \
86
- audioPath = os.path.join(self.cfg
87
- .audioPathAVA , self.cfg
88
- .evalDataType), \
89
- visualPath = os.path.join(self.cfg
90
- .visualPathAVA, self.cfg
91
- .evalDataType), \
92
- num_speakers = self.cfg.MODEL.NUM_SPEAKERS
93
- )
94
- valLoader = torch.utils.data.DataLoader(loader,
95
- batch_size=self.cfg.VAL.BATCH_SIZE,
96
- shuffle=False,
97
- pin_memory=True,
98
- num_workers=16)
99
-
100
- return valLoader
101
-
102
-
103
- def prepare_context_files(cfg):
104
- path = os.path.join(cfg.DATA.dataPathAVA, "csv")
105
- for phase in ["train", "val", "test"]:
106
- csv_f = f"{phase}_loader.csv"
107
- csv_orig = f"{phase}_orig.csv"
108
- entity_f = os.path.join(path, phase + "_entity.json")
109
- ts_f = os.path.join(path, phase + "_ts.json")
110
- if os.path.exists(entity_f) and os.path.exists(ts_f):
111
- continue
112
- orig_df = pandas.read_csv(os.path.join(path, csv_orig))
113
- entity_data = {}
114
- ts_to_entity = {}
115
-
116
- for index, row in orig_df.iterrows():
117
-
118
- entity_id = row['entity_id']
119
- video_id = row['video_id']
120
- if row['label'] == "SPEAKING_AUDIBLE":
121
- label = 1
122
- else:
123
- label = 0
124
- ts = float(row['frame_timestamp'])
125
- if video_id not in entity_data.keys():
126
- entity_data[video_id] = {}
127
- if entity_id not in entity_data[video_id].keys():
128
- entity_data[video_id][entity_id] = {}
129
- if ts not in entity_data[video_id][entity_id].keys():
130
- entity_data[video_id][entity_id][ts] = []
131
-
132
- entity_data[video_id][entity_id][ts] = label
133
-
134
- if video_id not in ts_to_entity.keys():
135
- ts_to_entity[video_id] = {}
136
- if ts not in ts_to_entity[video_id].keys():
137
- ts_to_entity[video_id][ts] = []
138
- ts_to_entity[video_id][ts].append(entity_id)
139
-
140
- with open(entity_f) as f:
141
- json.dump(entity_data, f)
142
-
143
- with open(ts_f) as f:
144
- json.dump(ts_to_entity, f)
145
-
146
-
147
- def main(gpu, world_size):
148
- # The structure of this code is learnt from https://github.com/clovaai/voxceleb_trainer
149
- cfg = bootstrap(print_cfg=False)
150
- rank = gpu
151
- dist.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=rank)
152
-
153
- make_deterministic(seed=int(cfg.SEED))
154
- torch.cuda.set_device(gpu)
155
- device = torch.device("cuda:{}".format(gpu))
156
-
157
- warnings.filterwarnings("ignore")
158
-
159
- cfg = init_args(cfg)
160
-
161
- data = DataPrep(cfg, world_size, rank)
162
-
163
- if cfg.downloadAVA == True:
164
- preprocess_AVA(cfg)
165
- quit()
166
-
167
- prepare_context_files(cfg)
168
-
169
- modelfiles = glob.glob('%s/model_0*.model' % cfg.modelSavePath)
170
- modelfiles.sort()
171
- if len(modelfiles) >= 1:
172
- print("Model %s loaded from previous state!" % modelfiles[-1])
173
- epoch = int(os.path.splitext(os.path.basename(modelfiles[-1]))[0][6:]) + 1
174
- s = loconet(cfg, rank, device)
175
- s.loadParameters(modelfiles[-1])
176
- else:
177
- epoch = 1
178
- s = loconet(cfg, rank, device)
179
-
180
- while (1):
181
- loss, lr = s.train_network(epoch=epoch, loader=data.train_dataloader())
182
-
183
- s.saveParameters(cfg.modelSavePath + "/model_%04d.model" % epoch)
184
-
185
- if epoch >= cfg.TRAIN.MAX_EPOCH:
186
- quit()
187
-
188
- epoch += 1
189
-
190
-
191
- if __name__ == '__main__':
192
-
193
- cfg = bootstrap()
194
- world_size = cfg.NUM_GPUS #
195
- os.environ['MASTER_ADDR'] = '127.0.0.1' #
196
- os.environ['MASTER_PORT'] = str(random.randint(4000, 8888)) #
197
- mp.spawn(main, nprocs=cfg.NUM_GPUS, args=(world_size,))