Pytorch—lightning初次使用

第一次使用lightning，之前看过很多关于lightning的讲解，直到看到这一篇才真的有点知道lightning的意义：从PyTorch到PyTorch Lightning —简要介绍 - 云+社区 - 腾讯云 (tencent.com)https://cloud.tencent.com/developer/article/1593703今天实际操作一下：

1.Data部分

Data部分目的是生成TrainDataLoader，需要数据准备，数据加载

1.1step部分：获取数据集 也就是pytorch里dataset部分

if stage == 'fit' or stage is None:

            self.trainset = video_datasets(self.tr_data_idx_dir, self.tr_data_frame_dir, transform = self.frame_trans )

1.2train_daatloader部分：加载数据集，返回训练集

def train_dataloader(self):
        return DataLoader(self.trainset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True)

2.模型部分

2.1 加载模型下面的方法为在另一个文件中获取模型

def load_model(self):
       
        self.model = AutoEncoderCov3DMem(self.kargs['ImgChnNum'], self.kargs['MemDim'], shrink_thres=self.kargs['ShrinkThres'])

下面的方法是在model部分之间定义模型

 def __init__(self):
    super(LightningMNISTClassifier, self).__init__()
 
    # mnist images are (1, 28, 28) (channels, width, height)
    self.layer_1 = torch.nn.Linear(28 * 28, 128)
    self.layer_2 = torch.nn.Linear(128, 256)
    self.layer_3 = torch.nn.Linear(256, 10)
 
  def forward(self, x):
      batch_size, channels, width, height = x.size()
 
      # (b, 1, 28, 28) -> (b, 1*28*28)
      x = x.view(batch_size, -1)
 
      # layer 1
      x = self.layer_1(x)
      x = torch.relu(x)
 
      # layer 2
      x = self.layer_2(x)
      x = torch.relu(x)
 
      # layer 3
      x = self.layer_3(x)
 
      # probability distribution over labels
      x = torch.log_softmax(x, dim=1)
 
      return x
 

2.2定义loss

def configure_loss(self):
        self.tr_recon_loss_func = nn.MSELoss()
        self.tr_entropy_loss_func = EntropyLossEncap()
        

2.3 定义优化器

def configure_optimizers(self):
        torch.optim.Adam(self.parameters(), lr = self.kargs.LR)

2.4 定义训练过程

def training_step(self, batch, batch_idx):
        img, labels, filename = batch
        out = self(img)
        loss = self.loss_function(out, labels)
        self.log('loss', loss, on_step=True, on_epoch=True, prog_bar=True)
        return loss

3.训练

定义trainer，fit进行训练

   data_module  = DInterface(**vars(args))
 
    model = MInterface(**vars(args))
    logger = TensorBoardLogger('tblog',name='testModel')
    trainer = Trainer.from_argparse_args(args,gpus = 3, logger = logger)
    trainer.fit(model, data_module)