基于PyTorch的「Keras」：除了核心逻辑通通都封装

# what to do in the training loopdef training_step(self, data_batch, batch_nb):# what to do in the validation loopdef validation_step(self, data_batch, batch_nb):
# how to aggregate validation_step outputsdef validation_end(self, outputs):
# and your dataloadersdef tng_dataloader():def val_dataloader():def test_dataloader():

import osimport torchfrom torch.nn import functional as Ffrom torch.utils.data import DataLoaderfrom torchvision.datasets import MNISTimport torchvision.transforms as transforms
import pytorch_lightning as ptl
class CoolModel(ptl.LightningModule):  def __init__(self):    super(CoolModel, self).__init__()    # not the best model...    self.l1 = torch.nn.Linear(28 * 28, 10)
  def forward(self, x):    return torch.relu(self.l1(x.view(x.size(0), -1)))    def my_loss(self, y_hat, y):    return F.cross_entropy(y_hat, y)      def training_step(self, batch, batch_nb):    x, y = batch    y_hat = self.forward(x)    return {'loss': self.my_loss(y_hat, y)}      def validation_step(self, batch, batch_nb):    x, y = batch    y_hat = self.forward(x)    return {'val_loss': self.my_loss(y_hat, y)}      def validation_end(self, outputs):    avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()    return {'avg_val_loss': avg_loss}    def configure_optimizers(self):    return [torch.optim.Adam(self.parameters(), lr=0.02)]
 @ptl.data_loader def tng_dataloader(self):   return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
 @ptl.data_loader def val_dataloader(self):   return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
 @ptl.data_loader def test_dataloader(self):   return DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)

from pytorch_lightning import Trainerfrom test_tube import Experiment

model = CoolModel()exp = Experiment(save_dir=os.getcwd())
# train on cpu using only 10% of the data (for demo purposes)trainer = Trainer(experiment=exp, max_nb_epochs=1, train_percent_check=0.1)
# train on 4 gpus# trainer = Trainer(experiment=exp, max_nb_epochs=1, gpus=[0, 1, 2, 3])# train on 32 gpus across 4 nodes (make sure to submit appropriate SLURM job)# trainer = Trainer(experiment=exp, max_nb_epochs=1, gpus=[0, 1, 2, 3, 4, 5, 6, 7], nb_gpu_nodes=4)# train (1 epoch only here for demo)trainer.fit(model)# view tensorflow logs print(f'View tensorboard logs by running\ntensorboard --logdir {os.getcwd()}')print('and going to http://localhost:6006 on your browser')

from test_tube import Experimentfrom pytorch-lightning import Trainer
exp = Experiment(save_dir = 『/some/path』)trainer = Trainer(experiment = exp)

tensorboard -logdir /some/path