Files
paddlepaddle--paddle/test/distributed_passes/model_zoo.py
T
2026-07-13 12:40:42 +08:00

110 lines
3.6 KiB
Python

# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import numpy as np
import paddle
from paddle import nn
from paddle.distributed import fleet
from paddle.vision.models import resnet50 as resnet
__all__ = [
'resnet_model',
]
def get_seed_from_env():
return int(os.environ.get("SEED", 0))
def resnet_model(
place, batch_size, image_shape=[3, 224, 224], num_classes=1000
):
image = paddle.static.data(
shape=[batch_size, *image_shape], dtype='float32', name='image'
)
label = paddle.static.data(
shape=[batch_size, 1], dtype='int64', name='label'
)
model = resnet(pretrained=False)
loss_fn = nn.loss.CrossEntropyLoss()
pred_out = model(image)
loss = loss_fn(pred_out, label)
optimizer = paddle.optimizer.Adam(learning_rate=1e-3)
dist_strategy = fleet.DistributedStrategy()
dist_strategy.fuse_all_reduce_ops = False
dist_strategy.without_graph_optimization = True
fleet.init(is_collective=True, strategy=dist_strategy)
optimizer = fleet.distributed_optimizer(optimizer)
optimizer.minimize(loss)
rank = paddle.distributed.get_rank()
def reader():
seed = get_seed_from_env()
np.random.seed(seed + rank)
for _ in range(10):
image_np = np.random.random(size=image.shape).astype('float32')
label_np = np.random.randint(
low=0, high=num_classes, size=label.shape
).astype('int64')
yield image_np, label_np
main_program = paddle.static.default_main_program()
startup_program = paddle.static.default_startup_program()
return main_program, startup_program, [image, label], [loss], reader
def simple_net(place, batch_size, image_shape=[784], num_classes=10):
image = paddle.static.data(
shape=[batch_size, *image_shape], dtype='float32', name='image'
)
label = paddle.static.data(
shape=[batch_size, 1], dtype='int64', name='label'
)
linears = [nn.Linear(784, 784) for _ in range(3)]
hidden = image
for linear in linears:
hidden = linear(hidden)
hidden = nn.ReLU()(hidden)
loss_fn = nn.loss.CrossEntropyLoss()
loss = loss_fn(hidden, label)
optimizer = paddle.optimizer.Adam(learning_rate=1e-3)
dist_strategy = fleet.DistributedStrategy()
dist_strategy.fuse_all_reduce_ops = False
dist_strategy.without_graph_optimization = True
fleet.init(is_collective=True, strategy=dist_strategy)
optimizer = fleet.distributed_optimizer(optimizer)
optimizer.minimize(loss)
rank = paddle.distributed.get_rank()
def reader():
seed = get_seed_from_env()
np.random.seed(seed + rank)
for _ in range(10):
image_np = np.random.random(size=image.shape).astype('float32')
label_np = np.random.randint(
low=0, high=num_classes, size=label.shape
).astype('int64')
yield image_np, label_np
main_program = paddle.static.default_main_program()
startup_program = paddle.static.default_startup_program()
return main_program, startup_program, [image, label], [loss], reader