Skip to content

Commit bde996c

Browse files
divideconceptdivideconcept
authored
0.9.10
1 parent ddbe9e0 commit bde996c

File tree

14 files changed

+3428
-3422
lines changed

14 files changed

+3428
-3422
lines changed

‎torchstudio/datasetanalyze.py‎

Lines changed: 117 additions & 117 deletions
Original file line numberDiff line numberDiff line change
@@ -1,117 +1,117 @@
1-
import sys
2-
3-
import torchstudio.tcpcodec as tc
4-
from torchstudio.modules import safe_exec
5-
import os
6-
import io
7-
from collections.abc import Iterable
8-
from tqdm.auto import tqdm
9-
import pickle
10-
11-
original_path=sys.path
12-
13-
app_socket = tc.connect()
14-
print("Analyze script connected\n", file=sys.stderr)
15-
while True:
16-
msg_type, msg_data = tc.recv_msg(app_socket)
17-
18-
if msg_type == 'SetAnalyzerCode':
19-
print("Setting analyzer code...\n", file=sys.stderr)
20-
analyzer = None
21-
analyzer_code = tc.decode_strings(msg_data)[0]
22-
error_msg, analyzer_env = safe_exec(analyzer_code, description='analyzer definition')
23-
if error_msg is not None or 'analyzer' not in analyzer_env:
24-
print("Unknown analyzer definition error" if error_msg is None else error_msg, file=sys.stderr)
25-
26-
if msg_type == 'StartAnalysisServer' and 'analyzer' in analyzer_env:
27-
print("Analyzing...\n", file=sys.stderr)
28-
29-
analysis_server, address = tc.generate_server()
30-
31-
if analyzer_env['analyzer'].train is None:
32-
request_msg='AnalysisServerRequestingAllSamples'
33-
elif analyzer_env['analyzer'].train==True:
34-
request_msg='AnalysisServerRequestingTrainingSamples'
35-
elif analyzer_env['analyzer'].train==False:
36-
request_msg='AnalysisServerRequestingValidationSamples'
37-
tc.send_msg(app_socket, request_msg, tc.encode_strings(address))
38-
dataset_socket=tc.start_server(analysis_server)
39-
40-
while True:
41-
dataset_msg_type, dataset_msg_data = tc.recv_msg(dataset_socket)
42-
43-
if dataset_msg_type == 'NumSamples':
44-
num_samples=tc.decode_ints(dataset_msg_data)[0]
45-
pbar=tqdm(total=num_samples, desc='Analyzing...', bar_format='{l_bar}{bar}| {remaining} left\n\n') #see https://github.com/tqdm/tqdm#parameters
46-
47-
if dataset_msg_type == 'InputTensorsID':
48-
input_tensors_id=tc.decode_ints(dataset_msg_data)
49-
50-
if dataset_msg_type == 'OutputTensorsID':
51-
output_tensors_id=tc.decode_ints(dataset_msg_data)
52-
53-
if dataset_msg_type == 'Labels':
54-
labels=tc.decode_strings(dataset_msg_data)
55-
56-
if dataset_msg_type == 'StartSending':
57-
error_msg, return_value = safe_exec(analyzer_env['analyzer'].start_analysis, (num_samples, input_tensors_id, output_tensors_id, labels), description='analyzer definition')
58-
if error_msg is not None:
59-
pbar.close()
60-
print(error_msg, file=sys.stderr)
61-
dataset_socket.close()
62-
analysis_server.close()
63-
break
64-
65-
if dataset_msg_type == 'TrainingSample':
66-
pbar.update(1)
67-
error_msg, return_value = safe_exec(analyzer_env['analyzer'].analyze_sample, (tc.decode_numpy_tensors(dataset_msg_data), True), description='analyzer definition')
68-
if error_msg is not None:
69-
pbar.close()
70-
print(error_msg, file=sys.stderr)
71-
dataset_socket.close()
72-
analysis_server.close()
73-
break
74-
75-
if dataset_msg_type == 'ValidationSample':
76-
pbar.update(1)
77-
error_msg, return_value = safe_exec(analyzer_env['analyzer'].analyze_sample, (tc.decode_numpy_tensors(dataset_msg_data), False), description='analyzer definition')
78-
if error_msg is not None:
79-
pbar.close()
80-
print(error_msg, file=sys.stderr)
81-
dataset_socket.close()
82-
analysis_server.close()
83-
break
84-
85-
if dataset_msg_type == 'DoneSending':
86-
pbar.close()
87-
error_msg, return_value = safe_exec(analyzer_env['analyzer'].finish_analysis, description='analyzer definition')
88-
tc.send_msg(dataset_socket, 'DoneReceiving')
89-
dataset_socket.close()
90-
analysis_server.close()
91-
if error_msg is not None:
92-
print(error_msg, file=sys.stderr)
93-
else:
94-
buffer=io.BytesIO()
95-
pickle.dump(analyzer_env['analyzer'].state_dict(), buffer)
96-
tc.send_msg(app_socket, 'AnalyzerState',buffer.getvalue())
97-
tc.send_msg(app_socket, 'AnalysisWeights',tc.encode_floats(analyzer_env['analyzer'].weights))
98-
print("Analysis complete")
99-
break
100-
101-
if msg_type == 'LoadAnalyzerState':
102-
if 'analyzer' in analyzer_env:
103-
buffer=io.BytesIO(msg_data)
104-
analyzer_env['analyzer'].load_state_dict(pickle.load(buffer))
105-
print("Analyzer state loaded")
106-
107-
if msg_type == 'RequestAnalysisReport':
108-
resolution = tc.decode_ints(msg_data)
109-
if 'analyzer' in analyzer_env:
110-
error_msg, return_value = safe_exec(analyzer_env['analyzer'].generate_report, (resolution[0:2],resolution[2]), description='analyzer definition')
111-
if error_msg is not None:
112-
print(error_msg, file=sys.stderr)
113-
if return_value is not None:
114-
tc.send_msg(app_socket, 'ReportImage', tc.encode_image(return_value))
115-
116-
if msg_type == 'Exit':
117-
break
1+
import sys
2+
3+
import torchstudio.tcpcodec as tc
4+
from torchstudio.modules import safe_exec
5+
import os
6+
import io
7+
from collections.abc import Iterable
8+
from tqdm.auto import tqdm
9+
import pickle
10+
11+
original_path=sys.path
12+
13+
app_socket = tc.connect()
14+
print("Analyze script connected\n", file=sys.stderr)
15+
while True:
16+
msg_type, msg_data = tc.recv_msg(app_socket)
17+
18+
if msg_type == 'SetAnalyzerCode':
19+
print("Setting analyzer code...\n", file=sys.stderr)
20+
analyzer = None
21+
analyzer_code = tc.decode_strings(msg_data)[0]
22+
error_msg, analyzer_env = safe_exec(analyzer_code, description='analyzer definition')
23+
if error_msg is not None or 'analyzer' not in analyzer_env:
24+
print("Unknown analyzer definition error" if error_msg is None else error_msg, file=sys.stderr)
25+
26+
if msg_type == 'StartAnalysisServer' and 'analyzer' in analyzer_env:
27+
print("Analyzing...\n", file=sys.stderr)
28+
29+
analysis_server, address = tc.generate_server()
30+
31+
if analyzer_env['analyzer'].train is None:
32+
request_msg='AnalysisServerRequestingAllSamples'
33+
elif analyzer_env['analyzer'].train==True:
34+
request_msg='AnalysisServerRequestingTrainingSamples'
35+
elif analyzer_env['analyzer'].train==False:
36+
request_msg='AnalysisServerRequestingValidationSamples'
37+
tc.send_msg(app_socket, request_msg, tc.encode_strings(address))
38+
dataset_socket=tc.start_server(analysis_server)
39+
40+
while True:
41+
dataset_msg_type, dataset_msg_data = tc.recv_msg(dataset_socket)
42+
43+
if dataset_msg_type == 'NumSamples':
44+
num_samples=tc.decode_ints(dataset_msg_data)[0]
45+
pbar=tqdm(total=num_samples, desc='Analyzing...', bar_format='{l_bar}{bar}| {remaining} left\n\n') #see https://github.com/tqdm/tqdm#parameters
46+
47+
if dataset_msg_type == 'InputTensorsID':
48+
input_tensors_id=tc.decode_ints(dataset_msg_data)
49+
50+
if dataset_msg_type == 'OutputTensorsID':
51+
output_tensors_id=tc.decode_ints(dataset_msg_data)
52+
53+
if dataset_msg_type == 'Labels':
54+
labels=tc.decode_strings(dataset_msg_data)
55+
56+
if dataset_msg_type == 'StartSending':
57+
error_msg, return_value = safe_exec(analyzer_env['analyzer'].start_analysis, (num_samples, input_tensors_id, output_tensors_id, labels), description='analyzer definition')
58+
if error_msg is not None:
59+
pbar.close()
60+
print(error_msg, file=sys.stderr)
61+
dataset_socket.close()
62+
analysis_server.close()
63+
break
64+
65+
if dataset_msg_type == 'TrainingSample':
66+
pbar.update(1)
67+
error_msg, return_value = safe_exec(analyzer_env['analyzer'].analyze_sample, (tc.decode_numpy_tensors(dataset_msg_data), True), description='analyzer definition')
68+
if error_msg is not None:
69+
pbar.close()
70+
print(error_msg, file=sys.stderr)
71+
dataset_socket.close()
72+
analysis_server.close()
73+
break
74+
75+
if dataset_msg_type == 'ValidationSample':
76+
pbar.update(1)
77+
error_msg, return_value = safe_exec(analyzer_env['analyzer'].analyze_sample, (tc.decode_numpy_tensors(dataset_msg_data), False), description='analyzer definition')
78+
if error_msg is not None:
79+
pbar.close()
80+
print(error_msg, file=sys.stderr)
81+
dataset_socket.close()
82+
analysis_server.close()
83+
break
84+
85+
if dataset_msg_type == 'DoneSending':
86+
pbar.close()
87+
error_msg, return_value = safe_exec(analyzer_env['analyzer'].finish_analysis, description='analyzer definition')
88+
tc.send_msg(dataset_socket, 'DoneReceiving')
89+
dataset_socket.close()
90+
analysis_server.close()
91+
if error_msg is not None:
92+
print(error_msg, file=sys.stderr)
93+
else:
94+
buffer=io.BytesIO()
95+
pickle.dump(analyzer_env['analyzer'].state_dict(), buffer)
96+
tc.send_msg(app_socket, 'AnalyzerState',buffer.getvalue())
97+
tc.send_msg(app_socket, 'AnalysisWeights',tc.encode_floats(analyzer_env['analyzer'].weights))
98+
print("Analysis complete")
99+
break
100+
101+
if msg_type == 'LoadAnalyzerState':
102+
if 'analyzer' in analyzer_env:
103+
buffer=io.BytesIO(msg_data)
104+
analyzer_env['analyzer'].load_state_dict(pickle.load(buffer))
105+
print("Analyzer state loaded")
106+
107+
if msg_type == 'RequestAnalysisReport':
108+
resolution = tc.decode_ints(msg_data)
109+
if 'analyzer' in analyzer_env:
110+
error_msg, return_value = safe_exec(analyzer_env['analyzer'].generate_report, (resolution[0:2],resolution[2]), description='analyzer definition')
111+
if error_msg is not None:
112+
print(error_msg, file=sys.stderr)
113+
if return_value is not None:
114+
tc.send_msg(app_socket, 'ReportImage', tc.encode_image(return_value))
115+
116+
if msg_type == 'Exit':
117+
break

0 commit comments

Comments
 (0)