时间戳

首页 > 解决方案 > 如何直接在 python 代码中“使用 -m 运行”而不是像“python -m ...”这样的命令？

python - 如何直接在 python 代码中“使用 -m 运行”而不是像“python -m ...”这样的命令？

问题描述

如何在 python 代码中直接使用“-m some-module”运行？

最近我了解到运行如下命令可以运行一些 torch.distributed python 模块：

“python -m torch.distributed.launch --nproc_per_node=2 SOME_TRAINING_SCRIPT.py”

========

如果我想通过执行脚本python SOME_TRAINING_SCRIPT.py（而不是python -m torch.distributed.launch --nproc_per_node=2 SOME_TRAINING_SCRIPT.py ）来实现通过 torch.distributed.launch 运行 SOME_TRAINING_SCRIPT.py ，应该添加哪些内容SOME_TRAINING_SCRIPT.py 来实现这个？

python文件如下：

进口火炬，操作系统
将 torch.nn 导入为 nn
从 torch.utils.data 导入数据集，DataLoader
从 torch.utils.data.distributed 导入 DistributedSampler
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"
os.environ['MASTER_ADDR'] = '127.0.0.1'
os.environ['MASTER_PORT'] = '29500'

torch.distributed.init_process_group(backend="gloo")

输入大小 = 5
输出大小 = 2
批量大小 = 30
数据大小 = 90

local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)

print("local_rank = ", local_rank)

类随机数据集（数据集）：
    def __init__(self, size, length):
        self.len = 长度
        self.data = torch.randn(长度，大小).to('cuda')

    def __getitem__(self, index):
        返回 self.data[索引]

    def __len__(self):
        返回 self.len


数据集 = RandomDataset(input_size, data_size)
rand_loader = DataLoader（数据集=数据集，
                         批量大小=批量大小，
                         采样器=分布式采样器（数据集））

类模型（nn.Module）：
    def __init__(self, input_size, output_size):
        超级（模型，自我）.__init__（）
        self.fc = nn.Linear(input_size, output_size)

    def 转发（自我，输入）：
        输出 = self.fc(输入)
        print("在模型中：输入尺寸", input.size(),
              "输出大小", output.size())
        返回输出


模型 = 模型（输入大小，输出大小）
模型.to（设备）

如果 torch.cuda.device_count() > 1：
    print("让我们使用", torch.cuda.device_count(), "GPUs!")
    模型 = torch.nn.parallel.DistributedDataParallel（模型，
                                                      device_ids=[local_rank],
                                                      输出设备=本地排名）

对于 rand_loader 中的数据：
    如果 torch.cuda.is_available():
        input_var = 数据
    别的：
        input_var = 数据
    输出 = 模型（输入变量）
    print("外部：输入大小", input_var.size(), "output_size", output.size())

标签： python

解决方案

如果您只想缩短训练脚本的运行命令，您可以使用bash脚本（Mac/Unix）简单地实现这一点。

Bash 脚本示例

创建一个名为run_distributed_training.sh内容的文件：

#!/bin/sh
python -m torch.distributed.launch --nproc_per_node=2 SOME_TRAINING_SCRIPT.py

然后你需要设置权限来执行这个：

chmod +x run_distributed_training.sh

最后，您可以像这样执行：

./run_distributed_training.sh

.bat您可以使用脚本为 Windows 实现相同的功能。

此处解释：为 python 脚本创建 BAT 文件

推荐阅读