Dataset

对于PyTorch而言data load的核心类是torch.utils.data.Dataloader，而对于创建一个新的Dataloder对象，最重要的参数就是dataset，一个torch.utils.data.Dataset对象。

PyTorch提供了两种风格的datasets，分别是map-style和iterable-style：

map-style datasets通过用户自行实现__getitem__()方法，可以做到随机读取。
iterable-style datasets通过实现__iter()__方法，可以以迭代的形式对数据进行读取。

这里需要考虑实验所使用的数据集形式。OpenFWI数据集的FlatVel-A类共有60个npy文件，每个npy文件包含500条数据。

这里，如果使用map-style datasets会存在无法避免的问题：由于Dataset支持随机读取，则实现__getitem__()时需要反复进行np.load()操作。这会导致读盘时间过长，拖慢训练速度。

因此，考虑使用IterableDataset类来规避这个问题：

python复制代码
1class Dataset(torch.utils.data.IterableDataset):
2    """
3    Dataset for InversionNet.
4    """
5
6    def __init__(self, root_dir, fid_list, num_samples_per_file=500):
7        """
8        Initialize dataset.
9
10        Args:
11            root_dir: root directory.
12            fid_list: list of npy file id.
13            num_samples_per_file: number of npy samples, which is 500 for OpenFWI.
14        """
15        super().__init__()
16        self.data_files = [
17            os.path.join(root_dir, "data", f"data{fid}.npy") for fid in fid_list
18        ]
19        self.label_files = [
20            os.path.join(root_dir, "model", f"model{fid}.npy") for fid in fid_list
21        ]
22        self.num_samples_per_file = num_samples_per_file
23
24    def __len__(self):
25        return len(self.data_files) * self.num_samples_per_file
26
27    def __iter__(self):
28        worker_info = torch.utils.data.get_worker_info()
29        if worker_info is not None:
30            num_workers = worker_info.num_workers
31            idx = worker_info.id
32            split_length = math.ceil(len(self.data_files) / num_workers)
33            data_range = range(
34                idx * split_length, min((idx + 1) * split_length, len(self.data_files))
35            )
36        else:
37            data_range = range(len(self.data_files))
38        for i in data_range:
39            data_file = np.load(self.data_files[i])
40            label_file = np.load(self.label_files[i])
41            for j in range(self.num_samples_per_file):
42                data, label = data_file[j], label_file[j]
43                yield data, label

需要注意的是，当设置num_workers为非零常数（使用multi-process data loading）后，系统会直接将整个dataset复制到每个进程中。因此在实现__iter__()时，需要手动做分片操作。

Model

之前跟Dive into deep learning时已经了解了大概，此处不展开。

python复制代码
1class InversionNet(nn.Module):
2    """
3    My InversionNet consisting of convolution block and deconvolution block.
4    """
5
6    def __init__(self, dim1=32, dim2=64, dim3=128, dim4=256, dim5=512, **kwargs):
7        """
8        Args:
9            dim1: Number of channels in the 1st layer
10            dim2: Number of channels in the 2nd layer
11            dim3: Number of channels in the 3rd layer
12            dim4: Number of channels in the 4th layer
13            dim5: Number of channels in the 5th layer
14        """
15        super().__init__()
16        // Init code
17
18    def forward(self, x):
19        // Forward code

Train

训练脚本主要由几个部分构成。

Parser

用于处理传入的命令参数，调用argparse库即可。

python复制代码
1def parse():
2    """
3    Create a new parser.
4    """
5    parser = argparse.ArgumentParser(
6        prog="Trainer", description="InversionNet Pytorch Trainer"
7    )
8    parser.add_argument(
9        "--batch-size",
10        type=int,
11        default=64,
12        metavar="N",
13        help="input batch size for training (default: 64)",
14    )
15    parser.add_argument(
16        "--train-size",
17        type=float,
18        default=0.8,
19        metavar="LR",
20        help="proportion for training dataset (default: 0.8)",
21    )
22    parser.add_argument(
23        "--epochs",
24        type=int,
25        default=1000,
26        metavar="N",
27        help="number of epochs to train (default: 1000)",
28    )
29    parser.add_argument(
30        "--lr",
31        type=float,
32        default=1e-2,
33        metavar="LR",
34        help="learning rate (default: 1e-2)",
35    )
36    parser.add_argument(
37        "--log-interval",
38        type=int,
39        default=1,
40        metavar="N",
41        help="how many batches to wait before logging training status",
42    )
43    parser.add_argument(
44        "--no-save-model",
45        action="store_true",
46        default=False,
47        help="do not save the current Model",
48    )
49    parser.add_argument(
50        "--seed", type=int, default=42, metavar="S", help="random seed (default: 42)"
51    )
52    parser.add_argument(
53        "--num-workers",
54        type=int,
55        default=0,
56        metavar="N",
57        help="number of processes if using multi-process data loading",
58    )
59    parser.add_argument(
60        "--no-cuda", action="store_true", default=False, help="disables CUDA training"
61    )
62    parser.add_argument(
63        "--no-mps",
64        action="store_true",
65        default=False,
66        help="disables macOS GPU training",
67    )
68    return parser

环境设置

设置训练所需的各种环境。

python复制代码
1torch.manual_seed(args.seed)
2use_cuda = not args.no_cuda and torch.cuda.is_available()
3use_mps = not args.no_mps and torch.backends.mps.is_available()
4if use_mps:
5    assert (
6        use_mps and args.num_workers == 0
7    ), "No support for multiprocess dataload using mps"
8kwargs = {
9    "batch_size": args.batch_size,
10    "num_workers": args.num_workers,
11    "pin_memory": True,
12}
13if use_cuda:
14    device = torch.device("cuda")
15elif use_mps:
16    device = torch.device("mps")
17else:
18    device = torch.device("cpu")

Dataloader

在主程序中创建对应的train_loader和test_loader，为了随机，对文件列表做了打乱处理。

python复制代码
1r = random.Random(args.seed)
2num_file = len(
3    [file for file in os.listdir("./data/data") if file.endswith(".npy")]
4)
5file_idx_list = [i for i in range(1, num_file + 1)]
6r.shuffle(file_idx_list)
7
8# Load dataset
9train_loader = torch.utils.data.DataLoader(
10        dataset=Dataset("./data", file_idx_list[: num_file * args.train_size]),
11        **kwargs,
12    )
13    test_loader = torch.utils.data.DataLoader(
14        dataset=Dataset("./data", file_idx_list[num_file * args.train_size :]),
15        **kwargs,
16    )

训练

工作站上搭载了两块显卡，因此创建DataParallel对象，调用两个GPU进行运算。

创建新网络、优化器等对象，开始训练：

python复制代码
1if use_cuda:
2    model = DataParallel(InversionNet())
3else:
4    model = InversionNet()
5model.to(device)
6optimizer = torch.optim.Adam(model.parameters(), args.lr)
7scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.6)
8
9loss, test_loss = [], []
10epochs = range(1, args.epochs + 1)
11for epoch in epochs:
12    loss.append(train(args, model, device, train_loader, optimizer, epoch))
13    test_loss.append(test(model, device, test_loader))
14    scheduler.step()

test()比较简单，主要逻辑只包含一个loss计算，重点分析train()函数。函数中包含下述语句：

python复制代码
1for batch_idx, (data, target) in enumerate(train_loader):
2  	// Code
3    return

当不设置num_workers时，len(enumerate(train_loader))是准确的，即math.ceil(len(train_loader.dataset) // batch_size)。然而，启用multi-process dataloading后，这个数值不再准确。这是因为len(enumerate(train_loader))仍然由上述的逻辑计算得到，而实际上batch的分配不再如此，具体逻辑如下：

根据num_workers数，和在IterableDataset处自定义的分片逻辑，得到每个进程需要处理的数据数量data_num。
对每个进程，分别按照batch_size进行数据获取。这样，每个进程的batch数就变成了math.ceil(data_num / batch_size)

根据上述原理，我们可以准确的计算出实际的batch数量：

python复制代码
1def cal_batch_num(args, dataloader):
2    if args.num_workers == 0:
3        batch_num = len(dataloader)
4    else:
5        batch_num = 0
6        for idx in range(args.num_workers):
7            split_length = math.ceil(len(dataloader.dataset) / args.num_workers)
8            data_num = min((idx + 1) * split_length, len(dataloader.dataset)) - idx * split_length
9            batch_num += math.ceil(data_num / args.batch_size)
10    return batch_num

至此，一个还算完整的炼丹代码就算写出来了，至于训练效果如何，那就看天意吧。

Reference

Comments