PyTorch最佳实践，怎样才能写出一手风格优美的代码(3)

请注意以下几点：

• 我们使用由「torchvision」包提供的预训练模型
• 我们将一个网络切分成三个模块，每个模块由预训练模型中的层组成
• 我们通过设置「requires_grad = False」来固定网络权重
• 我们返回一个带有三个模块输出的 list

4. 自定义损失函数

即使 PyTorch 已经具有了大量标准损失函数，你有时也可能需要创建自己的损失函数。为了做到这一点，你需要创建一个独立的「losses.py」文件，并且通过扩展「nn.Module」创建你的自定义损失函数：

class CustomLoss(torch.nn.Module): 
 
    def __init__(self): 
        super(CustomLoss,self).__init__() 
 
    def forward(self,x,y): 
        loss = torch.mean((x - y)**2) 
        return loss 

5. 训练模型的最佳代码结构

对于训练的最佳代码结构，我们需要使用以下两种模式：

• 使用 prefetch_generator 中的 BackgroundGenerator 来加载下一个批量数据
• 使用 tqdm 监控训练过程，并展示计算效率，这能帮助我们找到数据加载流程中的瓶颈

# import statements 
import torch 
import torch.nn as nn 
from torch.utils import data 
... 
 
# set flags / seeds 
torch.backends.cudnn.benchmark = True 
np.random.seed(1) 
torch.manual_seed(1) 
torch.cuda.manual_seed(1) 
... 
 
# Start with main code 
if __name__ == '__main__': 
    # argparse for additional flags for experiment 
    parser = argparse.ArgumentParser(description="Train a network for ...") 
    ... 
    opt = parser.parse_args()  
 
    # add code for datasets (we always use train and validation/ test set) 
    data_transforms = transforms.Compose([ 
        transforms.Resize((opt.img_size, opt.img_size)), 
        transforms.RandomHorizontalFlip(), 
        transforms.ToTensor(), 
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) 
    ]) 
 
    train_dataset = datasets.ImageFolder( 
        root=os.path.join(opt.path_to_data, "train"), 
        transform=data_transforms) 
    train_data_loader = data.DataLoader(train_dataset, ...) 
 
    test_dataset = datasets.ImageFolder( 
        root=os.path.join(opt.path_to_data, "test"), 
        transform=data_transforms) 
    test_data_loader = data.DataLoader(test_dataset ...) 
    ... 
 
    # instantiate network (which has been imported from *networks.py*) 
    net = MyNetwork(...) 
    ... 
 
    # create losses (criterion in pytorch) 
    criterion_L1 = torch.nn.L1Loss() 
    ... 
 
    # if running on GPU and we want to use cuda move model there 
    use_cuda = torch.cuda.is_available() 
    if use_cuda: 
        netnet = net.cuda() 
        ... 
 
    # create optimizers 
    optim = torch.optim.Adam(net.parameters(), lr=opt.lr) 
    ... 
 
    # load checkpoint if needed/ wanted 
    start_n_iter = 0 
    start_epoch = 0 
    if opt.resume: 
        ckpt = load_checkpoint(opt.path_to_checkpoint) # custom method for loading last checkpoint 
        net.load_state_dict(ckpt['net']) 
        start_epoch = ckpt['epoch'] 
        start_n_iter = ckpt['n_iter'] 
        optim.load_state_dict(ckpt['optim']) 
        print("last checkpoint restored") 
        ... 
 
    # if we want to run experiment on multiple GPUs we move the models there 
    net = torch.nn.DataParallel(net) 
    ... 
 
    # typically we use tensorboardX to keep track of experiments 
    writer = SummaryWriter(...) 
 
    # now we start the main loop 
    n_iter = start_n_iter 
    for epoch in range(start_epoch, opt.epochs): 
        # set models to train mode 
        net.train() 
        ... 
 
        # use prefetch_generator and tqdm for iterating through data 
        pbar = tqdm(enumerate(BackgroundGenerator(train_data_loader, ...)), 
                    total=len(train_data_loader)) 
        start_time = time.time() 
 
        # for loop going through dataset 
        for i, data in pbar: 
            # data preparation 
            img, label = data 
            if use_cuda: 
                imgimg = img.cuda() 
                labellabel = label.cuda() 
            ... 
 
            # It's very good practice to keep track of preparation time and computation time using tqdm to find any issues in your dataloader 
            prepare_time = start_time-time.time() 
 
            # forward and backward pass 
            optim.zero_grad() 
            ... 
            loss.backward() 
            optim.step() 
            ... 
 
            # udpate tensorboardX 
            writer.add_scalar(..., n_iter) 
            ... 
 
            # compute computation time and *compute_efficiency* 
            process_time = start_time-time.time()-prepare_time 
            pbar.set_description("Compute efficiency: {:.2f}, epoch: {}/{}:".format( 
                process_time/(process_time+prepare_time), epoch, opt.epochs)) 
            start_time = time.time() 
 
        # maybe do a test pass every x epochs 
        if epoch % x == x-1: 
            # bring models to evaluation mode 
            net.eval() 
            ... 
            #do some tests 
            pbar = tqdm(enumerate(BackgroundGenerator(test_data_loader, ...)), 
                    total=len(test_data_loader))  
            for i, data in pbar: 
                ... 
 
            # save checkpoint if needed 
            ... 

三、PyTorch 的多 GPU 训练

PyTorch 中有两种使用多 GPU 进行训练的模式。

根据我们的经验，这两种模式都是有效的。然而，第一种方法得到的结果更好、需要的代码更少。由于第二种方法中的 GPU 间的通信更少，似乎具有轻微的性能优势。

1. 对每个网络输入的 batch 进行切分

最常见的一种做法是直接将所有网络的输入切分为不同的批量数据，并分配给各个 GPU。

这样一来，在 1 个 GPU 上运行批量大小为 64 的模型，在 2 个 GPU 上运行时，每个 batch 的大小就变成了 32。这个过程可以使用「nn.DataParallel(model)」包装器自动完成。

2. 将所有网络打包到一个超级网络中，并对输入 batch 进行切分

这种模式不太常用。下面的代码仓库向大家展示了 Nvidia 实现的 pix2pixHD，它有这种方法的实现。

地址：https://github.com/NVIDIA/pix2pixHD

四、PyTorch 中该做和不该做的

1. 在「nn.Module」的「forward」方法中避免使用 Numpy 代码

Numpy 是在 CPU 上运行的，它比 torch 的代码运行得要慢一些。由于 torch 的开发思路与 numpy 相似，所以大多数 Numpy 中的函数已经在 PyTorch 中得到了支持。

2. 将「DataLoader」从主程序的代码中分离

载入数据的工作流程应该独立于你的主训练程序代码。PyTorch 使用「background」进程更加高效地载入数据，而不会干扰到主训练进程。

3. 不要在每一步中都记录结果

通常而言，我们要训练我们的模型好几千步。因此，为了减小计算开销，每隔 n 步对损失和其它的计算结果进行记录就足够了。尤其是，在训练过程中将中间结果保存成图像，这种开销是非常大的。

4. 使用命令行参数

（编辑：ASP站长网）