Pytorch_Lightning学习心得

Introduction

Pytorch_Lightning中与pytorch有一些格式上的差异，我认为Pytorch_Lightning的更为规范，让人读很繁琐的工程时能够直接快速定位，而且封装的更好，下面开始介绍一下相关的component。

Training

亮点：不需要手动将变量放入cuda；省去很多冗余操作，如zero_grad、backward、optimizer.step

在Pytorch中通常如下：

num_epoch = 5
for epoch in range(num_epoch):
    losses = list()
    accuracies = list()
    model.train()
    for batch in train_loader:
        x, y = batch

        # x: b x 1 x 28 x 28
        b = x. size(0)
        x = x.view(b,-1).cuda()

        # 1.forward
        1 = model (x)

        # 2.compute the objective function
        J = loss(1, y.cuda())
    
        # 3.cleaning the gradient
        model.zero_grad()                    # 很冗余
    
        # 4.accumlate the partial derivatives of J wrt params
        J.backward()                         # 很冗余

        # 5.step in the opposite direction of the gradient
        optimizer.step()                     # 很冗余

而在Pytorch_Lightning中，会简化这一步骤：

from pytorch_lightning.metrics.functional import accuracy

def training_step(self, batch, batch_idx):
    x, y = batch

    # x: b x 1 x 28 x 28
    b = x. size(0)
    x = x.view(b,-1)       # 不再需要指定cuda

    # 1.forward
    l = self(x)            # 因为这是在class中，用forward

    # 2.compute the objective function
    J = self.loss(1, y)    # 不再需要指定cuda，self.loss需要在class的__init__中定义好
    
    acc = accuracy(l, y)
    pbar = {'train_acc': acc}

    
    return {'loss': J, 'progress_bar': pbar}    # 以字典返回loss
    # return J               # 以值返回loss

Validation

基本上与Training类似

损失函数

在Pytorch中通常是先定义网络，专门写一个train函数，里面定义用到的loss，在Pytorch_Lightning中，是在网络中提前在网络的__init__中设置好loss，比如：

def __init__(self):
    super().__init__()
    self.loss = nn.CrossEntropyLoss()

优化器

• 单个优化器：当你只有一个优化器时，可以直接返回该优化器对象，PyTorch Lightning会自动处理，这里return [optimizer]也可以，但是一般单个不用列表：

def configure_optimizers(self):
    optimizer = optim.SGD(self.parameters(), lr=1e-2)
    return optimizer

• 多个优化器：当你有多个优化器时，需要返回一个包含这些优化器的列表：

def configure_optimizers(self):
    lr_g = self.lr_g_factor * self.learning_rate
    lr_d = self.learning_rate
    opt_ae = torch.optim.Adam(list(self.encoder.parameters()) +
                              list(self.decoder.parameters()) +
                              list(self.quant_conv.parameters()) +
                              list(self.post_quant_conv.parameters()),
                              lr=lr_g, betas=(0.5, 0.9))
    opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
                                lr=lr_d, betas=(0.5, 0.9))

    return [opt_ae, opt_disc]

• 多个优化器➕调度器

def configure_optimizers(self):
    lr_g = self.lr_g_factor * self.learning_rate
    lr_d = self.learning_rate
    opt_ae = torch.optim.Adam(list(self.encoder.parameters()) +
                              list(self.decoder.parameters()) +
                              list(self.quant_conv.parameters()) +
                              list(self.post_quant_conv.parameters()),
                              lr=lr_g, betas=(0.5, 0.9))
    opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
                                lr=lr_d, betas=(0.5, 0.9))

    scheduler_ae = torch.optim.lr_scheduler.StepLR(opt_ae, step_size=10, gamma=0.1)
    scheduler_disc = torch.optim.lr_scheduler.StepLR(opt_disc, step_size=10, gamma=0.1)

    return [opt_ae, opt_disc], [scheduler_ae, scheduler_disc]

原文地址：https://blog.csdn.net/weixin_45696917/article/details/140078749

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！