Merge pull request #5 from YiwenAI/main

update ray train

Author: YiwenAI

basic/test_ray_train.py

@@ -4,6 +4,10 @@
 from torchvision import datasets
 from torchvision.transforms import ToTensor
 
+from ray import train
+from ray.train.torch import TorchTrainer
+from ray.train import ScalingConfig
+
 
 def get_dataset():
     return datasets.FashionMNIST(
@@ -32,6 +36,7 @@ def forward(self, inputs):
         return logits
 
 
+# without distributed training, pure pytorch
 def train_func():
     num_epochs = 3
     batch_size = 64
@@ -54,4 +59,40 @@ def train_func():
         print(f"epoch: {epoch}, loss: {loss.item()}")
 
 
-train_func()
+# train_func()
+
+
+# distributed training
+def train_func_distributed():
+    num_epochs = 3
+    batch_size = 64
+
+    dataset = get_dataset()
+    dataloader = DataLoader(dataset, batch_size=batch_size)
+    dataloader = train.torch.prepare_data_loader(dataloader)
+
+    model = NeuralNetwork()
+    model = train.torch.prepare_model(model)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+
+    for epoch in range(num_epochs):
+        for inputs, labels in dataloader:
+            optimizer.zero_grad()
+            pred = model(inputs)
+            loss = criterion(pred, labels)
+            loss.backward()
+            optimizer.step()
+        print(f"epoch: {epoch}, loss: {loss.item()}")
+
+
+# For GPU Training, set `use_gpu` to True.
+use_gpu = False
+
+trainer = TorchTrainer(
+    train_func_distributed,
+    scaling_config=ScalingConfig(num_workers=4, use_gpu=use_gpu)
+)
+
+results = trainer.fit()