forked from mindspore-Ecosystem/mindspore
!9642 add gpu benchmark scripts for pynative mode
From: @chujinjin Reviewed-by: @jjfeing,@limingqi107 Signed-off-by: @limingqi107
This commit is contained in:
commit
651b1c3577
|
@ -813,6 +813,7 @@ class NASNetAMobile(nn.Cell):
|
|||
self.classifier = nn.Dense(in_channels=24*filters, out_channels=num_classes)
|
||||
self.shape = P.Shape()
|
||||
self.reshape = P.Reshape()
|
||||
self.avg_pool = nn.AvgPool2d(kernel_size=7, stride=1)
|
||||
self._initialize_weights()
|
||||
|
||||
def _initialize_weights(self):
|
||||
|
@ -867,7 +868,7 @@ class NASNetAMobile(nn.Cell):
|
|||
x_cell_15 = self.cell_15(x_cell_14, x_cell_13)
|
||||
|
||||
x_cell_15 = self.relu(x_cell_15)
|
||||
x_cell_15 = nn.AvgPool2d(F.shape(x_cell_15)[2:])(x_cell_15) # global average pool
|
||||
x_cell_15 = self.avg_pool(x_cell_15) # global average pool
|
||||
x_cell_15 = self.reshape(x_cell_15, (self.shape(x_cell_15)[0], -1,))
|
||||
x_cell_15 = self.dropout(x_cell_15)
|
||||
logits = self.classifier(x_cell_15)
|
||||
|
|
|
@ -39,6 +39,7 @@ parser.add_argument('--epoch_size', type=str, default="2", help='Epoch_size: def
|
|||
parser.add_argument('--print_per_steps', type=str, default="20", help='Print loss and time per steps: default 20')
|
||||
parser.add_argument('--run_distribute', type=ast.literal_eval, default=False, help='Run distribute')
|
||||
parser.add_argument('--dataset_path', type=str, default=None, help='Imagenet dataset path')
|
||||
parser.add_argument('--mode', type=str, default="GRAPH", choices=["GRAPH", "PYNATIVE"], help='Execute mode')
|
||||
parser.add_argument('--dtype', type=str, choices=["fp32", "fp16", "FP16", "FP32"], default="fp16",\
|
||||
help='Compute data type fp32 or fp16: default fp16')
|
||||
args_opt = parser.parse_args()
|
||||
|
@ -119,7 +120,11 @@ if __name__ == '__main__':
|
|||
compute_type = str(args_opt.dtype).lower()
|
||||
|
||||
# init context
|
||||
context.set_context(mode=context.GRAPH_MODE, device_target=dev, save_graphs=False)
|
||||
if args_opt.mode == "GRAPH":
|
||||
mode = context.GRAPH_MODE
|
||||
else:
|
||||
mode = context.PYNATIVE_MODE
|
||||
context.set_context(mode=mode, device_target=dev, save_graphs=False)
|
||||
if args_opt.run_distribute:
|
||||
init()
|
||||
context.set_auto_parallel_context(device_num=get_group_size(), parallel_mode=ParallelMode.DATA_PARALLEL,
|
||||
|
@ -174,10 +179,15 @@ if __name__ == '__main__':
|
|||
model = Model(net, loss_fn=loss, optimizer=opt, loss_scale_manager=loss_scale, metrics={'acc'},
|
||||
amp_level="O2", keep_batchnorm_fp32=False)
|
||||
# define callbacks
|
||||
if mode == context.PYNATIVE_MODE:
|
||||
print_per_steps = 1
|
||||
time_cb = MyTimeMonitor(total_batch, print_per_steps)
|
||||
loss_cb = LossMonitor()
|
||||
cb = [time_cb, loss_cb]
|
||||
|
||||
# train model
|
||||
print("========START RESNET50 GPU BENCHMARK========")
|
||||
model.train(int(epoch_size * step_size / print_per_steps), dataset, callbacks=cb, sink_size=print_per_steps)
|
||||
if mode == context.GRAPH_MODE:
|
||||
model.train(int(epoch_size * step_size / print_per_steps), dataset, callbacks=cb, sink_size=print_per_steps)
|
||||
else:
|
||||
model.train(epoch_size, dataset, callbacks=cb)
|
||||
|
|
Loading…
Reference in New Issue