【pytorch】代码整理

PyTorch的一个简单的网络

class ConvBlock(nn.Module):    
    def __init__(self):    
        super(ConvBlock, self).__init__()    
        block = [nn.Conv2d(...)]    
        block += [nn.ReLU()]    
        block += [nn.BatchNorm2d(...)]    
        self.block = nn.Sequential(*block)    
    def forward(self, x):    
        return self.block(x)    
class SimpleNetwork(nn.Module):    
    def __init__(self, num_resnet_blocks=6):    
        super(SimpleNetwork, self).__init__()    
        # here we add the individual layers    
        layers = [ConvBlock(...)]    
        for i in range(num_resnet_blocks):    
            layers += [ResBlock(...)]    
        self.net = nn.Sequential(*layers)    
    def forward(self, x):    
        return self.net(x)

请注意以下几点：

•     我们重用简单的循环构建块，如ConvBlock，它由相同的循环模式(卷积、激活、归一化)组成，并将它们放入单独的nn.Module中；
•     我们建立一个所需层的列表，最后使用nn.Sequential()将它们转换成一个模型。我们在list对象之前使用*操作符来展开它。
•     在前向传递中，我们只是通过模型运行输入

在PyTorch中使用残差链接的网络

    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):    
        super(ResnetBlock, self).__init__()    
        self.conv_block = self.build_conv_block(...)    
    def build_conv_block(self, ...):    
        conv_block = []    
        conv_block += [nn.Conv2d(...),    
                       norm_layer(...),    
                       nn.ReLU()]    
        if use_dropout:    
            conv_block += [nn.Dropout(...)]    
        conv_block += [nn.Conv2d(...),    
                       norm_layer(...)]    
        return nn.Sequential(*conv_block)    
    def forward(self, x):    
        out = x + self.conv_block(x)    
        return out

在PyTorch使用多个输出的网络

class Vgg19(torch.nn.Module):    
  def __init__(self, requires_grad=False):    
    super(Vgg19, self).__init__()    
    vgg_pretrained_features = models.vgg19(pretrained=True).features    
    self.slice1 = torch.nn.Sequential()    
    self.slice2 = torch.nn.Sequential()    
    self.slice3 = torch.nn.Sequential()    
    for x in range(7):    
        self.slice1.add_module(str(x), vgg_pretrained_features[x])    
    for x in range(7, 21):    
        self.slice2.add_module(str(x), vgg_pretrained_features[x])    
    for x in range(21, 30):    
        self.slice3.add_module(str(x), vgg_pretrained_features[x])    
    if not requires_grad:    
        for param in self.parameters():    
            param.requires_grad = False    
  def forward(self, x):    
    h_relu1 = self.slice1(x)    
    h_relu2 = self.slice2(h_relu1)            
    h_relu3 = self.slice3(h_relu2)            
    out = [h_relu1, h_relu2, h_relu3]    
    return out

自定义Loss

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

训练模型的推荐代码结构

# 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:    
        net = 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:    
                img = img.cuda()    
                label = 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之间的通信更少，第二种方法似乎具有轻微的性能优势。

分割每个网络的batch

最常见的一种方法是简单地将所有“网络”的batch分配给各个gpu。

因此，如果一个模型运行在一个批处理大小为64的GPU上，那么它将运行在两个GPU上，每个GPU的批处理大小为32。这可以通过使用nn.DataParallel(model)自动完成。

将所有的网络打包进一个super网络，并把输入batch分割

这种模式不太常用。实现这种方法的repository在pix2pixHD implementation by Nvidia

该做的和不该做的

避免在nn.Module的forward方法找那个使用Numpy代码

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

从main代码中分离DataLoader

数据加载管道应该独立于你的主训练代码。PyTorch使用后台来更有效地加载数据，并且不会干扰主训练过程。

不要在每一次迭代中打印日志结果

通常我们训练我们的模型数千个迭代。因此，每n步记录损失和其他结果就足以减少开销。特别是，在训练过程中，将中间结果保存为图像可能非常耗时。

使用命令行参数

使用命令行参数在代码执行期间设置参数(批处理大小、学习率等)非常方便。跟踪实验参数的一个简单方法是打印从parse_args接收到的字典：

...    
# saves arguments to config.txt file    
opt = parser.parse_args()    
with open("config.txt", "w") as f:    
    f.write(opt.__str__())    
...

可能的话，使用.detach()将张量从图中释放出来

PyTorch跟踪所有涉及张量的操作，以实现自动微分。使用.detach()防止记录不必要的操作。

使用.item()打印标量数据

你可以直接打印变量，但是建议使用variable.detach()或variable.item()。在早期的PyTorch版本< 0.4中，必须使用.data访问一个变量的张量。

在nn.Module中使用函数调用而不是直接用forward

下面这两种方式是不一样的：

output = self.net.forward(input)    
# they are not equal!    
output = self.net(input)

 

 

 

 

 

---

 

参考资料：

1. https://blog.csdn.net/u011984148/article/details/99440021