22
22
parser = argparse .ArgumentParser (description = __doc__ )
23
23
add_arg = functools .partial (add_arguments , argparser = parser )
24
24
# yapf: disable
25
- add_arg ('batch_size' , int , 64 , "Minibatch size." )
25
+ add_arg ('batch_size' , int , 256 , "Minibatch size." )
26
26
add_arg ('use_gpu' , bool , True , "Whether to use GPU or not." )
27
27
add_arg ('save_inference' , bool , False , "Whether to save inference model." )
28
28
add_arg ('total_images' , int , 1281167 , "Training image number." )
45
45
46
46
47
47
def piecewise_decay (args ):
48
- if args .use_gpu :
49
- devices_num = paddle .fluid .core .get_cuda_device_count ()
50
- else :
51
- devices_num = int (os .environ .get ('CPU_NUM' , 1 ))
52
- step = int (
53
- math .ceil (float (args .total_images ) / args .batch_size ) / devices_num )
48
+ step = int (math .ceil (float (args .total_images ) / args .batch_size ))
54
49
bd = [step * e for e in args .step_epochs ]
55
50
lr = [args .lr * (0.1 ** i ) for i in range (len (bd ) + 1 )]
56
51
learning_rate = paddle .optimizer .lr .PiecewiseDecay (
@@ -63,12 +58,7 @@ def piecewise_decay(args):
63
58
64
59
65
60
def cosine_decay (args ):
66
- if args .use_gpu :
67
- devices_num = paddle .fluid .core .get_cuda_device_count ()
68
- else :
69
- devices_num = int (os .environ .get ('CPU_NUM' , 1 ))
70
- step = int (
71
- math .ceil (float (args .total_images ) / args .batch_size ) / devices_num )
61
+ step = int (math .ceil (float (args .total_images ) / args .batch_size ))
72
62
learning_rate = paddle .optimizer .lr .CosineAnnealingDecay (
73
63
learning_rate = args .lr , T_max = step * args .num_epochs , verbose = False )
74
64
optimizer = paddle .optimizer .Momentum (
@@ -108,7 +98,10 @@ def compress(args):
108
98
places = paddle .static .cuda_places (
109
99
) if args .use_gpu else paddle .static .cpu_places ()
110
100
place = places [0 ]
111
-
101
+ if args .use_gpu :
102
+ devices_num = paddle .fluid .core .get_cuda_device_count ()
103
+ else :
104
+ devices_num = int (os .environ .get ('CPU_NUM' , 1 ))
112
105
with paddle .static .program_guard (student_program , s_startup ):
113
106
with paddle .fluid .unique_name .guard ():
114
107
image = paddle .static .data (
@@ -120,7 +113,7 @@ def compress(args):
120
113
places = places ,
121
114
feed_list = [image , label ],
122
115
drop_last = True ,
123
- batch_size = args .batch_size ,
116
+ batch_size = int ( args .batch_size / devices_num ) ,
124
117
return_list = False ,
125
118
shuffle = True ,
126
119
use_shared_memory = True ,
0 commit comments