@@ -90,7 +90,7 @@ def train_program(is_sparse):
90
90
return avg_cost
91
91
92
92
93
- def train (use_cuda , train_program , save_path ):
93
+ def train (use_cuda , train_program , save_dirname ):
94
94
train_reader = paddle .batch (
95
95
paddle .dataset .imikolov .train (word_dict , N ), BATCH_SIZE )
96
96
test_reader = paddle .batch (
@@ -99,50 +99,69 @@ def train(use_cuda, train_program, save_path):
99
99
place = fluid .CUDAPlace (0 ) if use_cuda else fluid .CPUPlace ()
100
100
101
101
def event_handler (event ):
102
- if isinstance (event , fluid .EndEpochEvent ):
103
- outs = trainer .test (reader = test_reader )
102
+ if isinstance (event , fluid .EndStepEvent ):
103
+ outs = trainer .test (
104
+ reader = test_reader ,
105
+ feed_order = ['firstw' , 'secondw' , 'thirdw' , 'forthw' , 'nextw' ])
104
106
avg_cost = outs [0 ]
105
107
print ("loss= " , avg_cost )
106
108
107
- if avg_cost < 5.0 :
108
- trainer .save_params (save_path )
109
- return
109
+ if avg_cost < 10.0 :
110
+ trainer .save_params (save_dirname )
111
+ trainer .stop ()
112
+
110
113
if math .isnan (avg_cost ):
111
114
sys .exit ("got NaN loss, training failed." )
112
115
113
116
trainer = fluid .Trainer (
114
- train_program , fluid .optimizer .SGD (learning_rate = 0.001 ), place = place )
117
+ train_func = train_program ,
118
+ optimizer = fluid .optimizer .SGD (learning_rate = 0.001 ),
119
+ place = place )
120
+
115
121
trainer .train (
116
- reader = train_reader , num_epochs = 1 , event_handler = event_handler )
122
+ reader = train_reader ,
123
+ num_epochs = 1 ,
124
+ event_handler = event_handler ,
125
+ feed_order = ['firstw' , 'secondw' , 'thirdw' , 'forthw' , 'nextw' ])
117
126
118
127
119
- def infer (use_cuda , inference_program , save_path ):
128
+ def infer (use_cuda , inference_program , save_dirname = None ):
120
129
place = fluid .CUDAPlace (0 ) if use_cuda else fluid .CPUPlace ()
121
130
inferencer = fluid .Inferencer (
122
- infer_func = inference_program , param_path = save_path , place = place )
131
+ infer_func = inference_program , param_path = save_dirname , place = place )
123
132
124
133
lod = [0 , 1 ]
125
134
first_word = create_random_lodtensor (lod , place , low = 0 , high = dict_size - 1 )
126
135
second_word = create_random_lodtensor (lod , place , low = 0 , high = dict_size - 1 )
127
136
third_word = create_random_lodtensor (lod , place , low = 0 , high = dict_size - 1 )
128
137
fourth_word = create_random_lodtensor (lod , place , low = 0 , high = dict_size - 1 )
129
138
130
- result = inferencer .infer ({
131
- 'firstw' : first_word ,
132
- 'secondw' : second_word ,
133
- 'thirdw' : third_word ,
134
- 'forthw' : fourth_word
135
- })
139
+ result = inferencer .infer (
140
+ {
141
+ 'firstw' : first_word ,
142
+ 'secondw' : second_word ,
143
+ 'thirdw' : third_word ,
144
+ 'forthw' : fourth_word
145
+ },
146
+ return_numpy = False )
136
147
print (np .array (result [0 ]))
137
148
138
149
139
150
def main (use_cuda , is_sparse ):
140
151
if use_cuda and not fluid .core .is_compiled_with_cuda ():
141
152
return
142
153
143
- save_path = "word2vec.params"
144
- train (use_cuda , partial (train_program , is_sparse ), save_path )
145
- infer (use_cuda , partial (inference_program , is_sparse ), save_path )
154
+ save_path = "word2vec.inference.model"
155
+
156
+ train (
157
+ use_cuda = use_cuda ,
158
+ train_program = partial (train_program , is_sparse ),
159
+ save_dirname = save_path )
160
+
161
+ infer (
162
+ use_cuda = use_cuda ,
163
+ inference_program = partial (inference_program , is_sparse ),
164
+ save_dirname = save_path )
146
165
147
166
148
167
if __name__ == '__main__' :
0 commit comments