chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
+161
@@ -0,0 +1,161 @@
|
||||
# The PyTorch data transfer benchmark script.
|
||||
import argparse
|
||||
import warnings
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
import ray.train as train
|
||||
from ray.train import ScalingConfig
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
|
||||
class Net(nn.Module):
|
||||
def __init__(self, in_d, hidden):
|
||||
# output dim = 1
|
||||
super(Net, self).__init__()
|
||||
dims = [in_d] + hidden + [1]
|
||||
self.layers = nn.ModuleList(
|
||||
[nn.Linear(dims[i - 1], dims[i]) for i in range(len(dims))]
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
for layer in self.layers:
|
||||
x = layer(x)
|
||||
return x
|
||||
|
||||
|
||||
class BenchmarkDataset(torch.utils.data.Dataset):
|
||||
"""Create a naive dataset for the benchmark"""
|
||||
|
||||
def __init__(self, dim, size=1000):
|
||||
self.x = torch.from_numpy(np.random.normal(size=(size, dim))).float()
|
||||
self.y = torch.from_numpy(np.random.normal(size=(size, 1))).float()
|
||||
self.size = size
|
||||
|
||||
def __getitem__(self, index):
|
||||
return self.x[index, None], self.y[index, None]
|
||||
|
||||
def __len__(self):
|
||||
return self.size
|
||||
|
||||
|
||||
def train_epoch(epoch, dataloader, model, loss_fn, optimizer):
|
||||
if train.get_context().get_world_size() > 1:
|
||||
dataloader.sampler.set_epoch(epoch)
|
||||
|
||||
for X, y in dataloader:
|
||||
# Compute prediction error
|
||||
pred = model(X)
|
||||
loss = loss_fn(pred, y)
|
||||
|
||||
# Backpropagation
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
|
||||
def train_func(config):
|
||||
data_size = config.get("data_size", 4096 * 50)
|
||||
batch_size = config.get("batch_size", 4096)
|
||||
hidden_size = config.get("hidden_size", 1)
|
||||
use_auto_transfer = config.get("use_auto_transfer", False)
|
||||
lr = config.get("lr", 1e-2)
|
||||
epochs = config.get("epochs", 10)
|
||||
|
||||
train_dataset = BenchmarkDataset(4096, size=data_size)
|
||||
train_loader = torch.utils.data.DataLoader(
|
||||
train_dataset, batch_size=batch_size, shuffle=True
|
||||
)
|
||||
|
||||
train_loader = train.torch.prepare_data_loader(
|
||||
data_loader=train_loader, move_to_device=True, auto_transfer=use_auto_transfer
|
||||
)
|
||||
|
||||
model = Net(in_d=4096, hidden=[4096] * hidden_size)
|
||||
model = train.torch.prepare_model(model)
|
||||
|
||||
loss_fn = nn.MSELoss()
|
||||
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
|
||||
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
|
||||
choice = "with" if use_auto_transfer else "without"
|
||||
print(f"Starting the torch data prefetch benchmark {choice} auto pipeline...")
|
||||
|
||||
torch.cuda.synchronize()
|
||||
start.record()
|
||||
for epoch in range(epochs):
|
||||
train_epoch(epoch, train_loader, model, loss_fn, optimizer)
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
print(
|
||||
f"Finished the torch data prefetch benchmark {choice} "
|
||||
f"auto pipeline: {start.elapsed_time(end)} ms."
|
||||
)
|
||||
|
||||
return "Experiment done."
|
||||
|
||||
|
||||
def train_linear(num_workers=1, num_hidden_layers=1, use_auto_transfer=True, epochs=3):
|
||||
config = {
|
||||
"lr": 1e-2,
|
||||
"hidden_size": num_hidden_layers,
|
||||
"batch_size": 4096,
|
||||
"epochs": epochs,
|
||||
"use_auto_transfer": use_auto_transfer,
|
||||
}
|
||||
trainer = TorchTrainer(
|
||||
train_func,
|
||||
train_loop_config=config,
|
||||
scaling_config=ScalingConfig(use_gpu=True, num_workers=num_workers),
|
||||
)
|
||||
results = trainer.fit()
|
||||
|
||||
print(results.metrics)
|
||||
return results
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--address", required=False, type=str, help="the address to use for Ray"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--epochs", type=int, default=1, help="Number of epochs to train for."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num_hidden_layers",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of epochs to train for.",
|
||||
)
|
||||
|
||||
args, _ = parser.parse_known_args()
|
||||
|
||||
import ray
|
||||
|
||||
ray.init(address=args.address)
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
warnings.warn("GPU is not available. Skip the test using auto pipeline.")
|
||||
else:
|
||||
train_linear(
|
||||
num_workers=1,
|
||||
num_hidden_layers=args.num_hidden_layers,
|
||||
use_auto_transfer=True,
|
||||
epochs=args.epochs,
|
||||
)
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
train_linear(
|
||||
num_workers=1,
|
||||
num_hidden_layers=args.num_hidden_layers,
|
||||
use_auto_transfer=False,
|
||||
epochs=args.epochs,
|
||||
)
|
||||
|
||||
ray.shutdown()
|
||||
@@ -0,0 +1,154 @@
|
||||
import os
|
||||
from typing import Dict
|
||||
|
||||
import torch
|
||||
from filelock import FileLock
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader
|
||||
from torchvision import datasets, transforms
|
||||
from torchvision.transforms import Normalize, ToTensor
|
||||
from tqdm import tqdm
|
||||
|
||||
import ray.train
|
||||
from ray.train import ScalingConfig
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
|
||||
def get_dataloaders(batch_size):
|
||||
# Transform to normalize the input images
|
||||
transform = transforms.Compose([ToTensor(), Normalize((0.28604,), (0.32025,))])
|
||||
|
||||
with FileLock(os.path.expanduser("~/data.lock")):
|
||||
# Download training data from open datasets
|
||||
training_data = datasets.FashionMNIST(
|
||||
root="~/data",
|
||||
train=True,
|
||||
download=True,
|
||||
transform=transform,
|
||||
)
|
||||
|
||||
# Download test data from open datasets
|
||||
test_data = datasets.FashionMNIST(
|
||||
root="~/data",
|
||||
train=False,
|
||||
download=True,
|
||||
transform=transform,
|
||||
)
|
||||
|
||||
# Create data loaders
|
||||
train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True)
|
||||
test_dataloader = DataLoader(test_data, batch_size=batch_size)
|
||||
|
||||
return train_dataloader, test_dataloader
|
||||
|
||||
|
||||
# Model Definition
|
||||
class NeuralNetwork(nn.Module):
|
||||
def __init__(self):
|
||||
super(NeuralNetwork, self).__init__()
|
||||
self.flatten = nn.Flatten()
|
||||
self.linear_relu_stack = nn.Sequential(
|
||||
nn.Linear(28 * 28, 512),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.25),
|
||||
nn.Linear(512, 512),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.25),
|
||||
nn.Linear(512, 10),
|
||||
nn.ReLU(),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.flatten(x)
|
||||
logits = self.linear_relu_stack(x)
|
||||
return logits
|
||||
|
||||
|
||||
def train_func_per_worker(config: Dict):
|
||||
ray.train.torch.enable_reproducibility()
|
||||
|
||||
lr = config["lr"]
|
||||
epochs = config["epochs"]
|
||||
batch_size = config["batch_size_per_worker"]
|
||||
|
||||
# Get dataloaders inside the worker training function
|
||||
train_dataloader, test_dataloader = get_dataloaders(batch_size=batch_size)
|
||||
|
||||
# [1] Prepare Dataloader for distributed training
|
||||
# Shard the datasets among workers and move batches to the correct device
|
||||
# =======================================================================
|
||||
train_dataloader = ray.train.torch.prepare_data_loader(train_dataloader)
|
||||
test_dataloader = ray.train.torch.prepare_data_loader(test_dataloader)
|
||||
|
||||
model = NeuralNetwork()
|
||||
|
||||
# [2] Prepare and wrap your model with DistributedDataParallel
|
||||
# Move the model to the correct GPU/CPU device
|
||||
# ============================================================
|
||||
model = ray.train.torch.prepare_model(model)
|
||||
|
||||
loss_fn = nn.CrossEntropyLoss()
|
||||
optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)
|
||||
|
||||
# Model training loop
|
||||
for epoch in range(epochs):
|
||||
if ray.train.get_context().get_world_size() > 1:
|
||||
# Required for the distributed sampler to shuffle properly across epochs.
|
||||
train_dataloader.sampler.set_epoch(epoch)
|
||||
|
||||
model.train()
|
||||
for X, y in tqdm(train_dataloader, desc=f"Train Epoch {epoch}"):
|
||||
pred = model(X)
|
||||
loss = loss_fn(pred, y)
|
||||
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
model.eval()
|
||||
test_loss, num_correct, num_total = 0, 0, 0
|
||||
with torch.no_grad():
|
||||
for X, y in tqdm(test_dataloader, desc=f"Test Epoch {epoch}"):
|
||||
pred = model(X)
|
||||
loss = loss_fn(pred, y)
|
||||
|
||||
test_loss += loss.item()
|
||||
num_total += y.shape[0]
|
||||
num_correct += (pred.argmax(1) == y).sum().item()
|
||||
|
||||
test_loss /= len(test_dataloader)
|
||||
accuracy = num_correct / num_total
|
||||
|
||||
# [3] Report metrics to Ray Train
|
||||
# ===============================
|
||||
ray.train.report(metrics={"loss": test_loss, "accuracy": accuracy})
|
||||
|
||||
|
||||
def train_fashion_mnist(num_workers=2, use_gpu=False):
|
||||
global_batch_size = 32
|
||||
|
||||
train_config = {
|
||||
"lr": 1e-3,
|
||||
"epochs": 10,
|
||||
"batch_size_per_worker": global_batch_size // num_workers,
|
||||
}
|
||||
|
||||
# Configure computation resources
|
||||
scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu)
|
||||
|
||||
# Initialize a Ray TorchTrainer
|
||||
trainer = TorchTrainer(
|
||||
train_loop_per_worker=train_func_per_worker,
|
||||
train_loop_config=train_config,
|
||||
scaling_config=scaling_config,
|
||||
)
|
||||
|
||||
# [4] Start distributed training
|
||||
# Run `train_func_per_worker` on all workers
|
||||
# =============================================
|
||||
result = trainer.fit()
|
||||
print(f"Training result: {result}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
train_fashion_mnist(num_workers=4, use_gpu=True)
|
||||
@@ -0,0 +1,147 @@
|
||||
import argparse
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
import ray.train as train
|
||||
from ray.train import Checkpoint, RunConfig, ScalingConfig
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
|
||||
class LinearDataset(torch.utils.data.Dataset):
|
||||
"""y = a * x + b"""
|
||||
|
||||
def __init__(self, a, b, size=1000):
|
||||
x = np.arange(0, 10, 10 / size, dtype=np.float32)
|
||||
self.x = torch.from_numpy(x)
|
||||
self.y = torch.from_numpy(a * x + b)
|
||||
|
||||
def __getitem__(self, index):
|
||||
return self.x[index, None], self.y[index, None]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.x)
|
||||
|
||||
|
||||
def train_epoch(epoch, dataloader, model, loss_fn, optimizer):
|
||||
if train.get_context().get_world_size() > 1:
|
||||
dataloader.sampler.set_epoch(epoch)
|
||||
|
||||
for X, y in dataloader:
|
||||
# Compute prediction error
|
||||
pred = model(X)
|
||||
loss = loss_fn(pred, y)
|
||||
|
||||
# Backpropagation
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
|
||||
def validate_epoch(dataloader, model, loss_fn):
|
||||
num_batches = len(dataloader)
|
||||
model.eval()
|
||||
loss = 0
|
||||
with torch.no_grad():
|
||||
for X, y in dataloader:
|
||||
pred = model(X)
|
||||
loss += loss_fn(pred, y).item()
|
||||
loss /= num_batches
|
||||
import copy
|
||||
|
||||
model_copy = copy.deepcopy(model)
|
||||
return model_copy.cpu().state_dict(), loss
|
||||
|
||||
|
||||
def train_func(config):
|
||||
data_size = config.get("data_size", 1000)
|
||||
val_size = config.get("val_size", 400)
|
||||
batch_size = config.get("batch_size", 32)
|
||||
hidden_size = config.get("hidden_size", 1)
|
||||
lr = config.get("lr", 1e-2)
|
||||
epochs = config.get("epochs", 3)
|
||||
|
||||
train_dataset = LinearDataset(2, 5, size=data_size)
|
||||
val_dataset = LinearDataset(2, 5, size=val_size)
|
||||
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size)
|
||||
validation_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size)
|
||||
|
||||
train_loader = train.torch.prepare_data_loader(train_loader)
|
||||
validation_loader = train.torch.prepare_data_loader(validation_loader)
|
||||
|
||||
model = nn.Linear(1, hidden_size)
|
||||
model = train.torch.prepare_model(model)
|
||||
|
||||
loss_fn = nn.MSELoss()
|
||||
|
||||
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
|
||||
|
||||
results = []
|
||||
for epoch in range(epochs):
|
||||
train_epoch(epoch, train_loader, model, loss_fn, optimizer)
|
||||
state_dict, loss = validate_epoch(validation_loader, model, loss_fn)
|
||||
result = dict(loss=loss)
|
||||
results.append(result)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
torch.save(state_dict, os.path.join(tmpdir, "model.pt"))
|
||||
train.report(result, checkpoint=Checkpoint.from_directory(tmpdir))
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def train_linear(num_workers=2, use_gpu=False, epochs=3, storage_path=None):
|
||||
config = {"lr": 1e-2, "hidden_size": 1, "batch_size": 4, "epochs": epochs}
|
||||
trainer = TorchTrainer(
|
||||
train_loop_per_worker=train_func,
|
||||
train_loop_config=config,
|
||||
scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
|
||||
run_config=RunConfig(storage_path=storage_path),
|
||||
)
|
||||
result = trainer.fit()
|
||||
|
||||
print(result.metrics)
|
||||
return result.metrics
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--address", required=False, type=str, help="the address to use for Ray"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-workers",
|
||||
"-n",
|
||||
type=int,
|
||||
default=2,
|
||||
help="Sets number of workers for training.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use-gpu", action="store_true", help="Whether to use GPU for training."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--epochs", type=int, default=3, help="Number of epochs to train for."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--smoke-test",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Finish quickly for testing.",
|
||||
)
|
||||
|
||||
args, _ = parser.parse_known_args()
|
||||
|
||||
import ray
|
||||
|
||||
if args.smoke_test:
|
||||
# 2 workers + 1 for trainer.
|
||||
ray.init(num_cpus=3)
|
||||
train_linear()
|
||||
else:
|
||||
ray.init(address=args.address)
|
||||
train_linear(
|
||||
num_workers=args.num_workers, use_gpu=args.use_gpu, epochs=args.epochs
|
||||
)
|
||||
@@ -0,0 +1,110 @@
|
||||
# ruff: noqa
|
||||
# fmt: off
|
||||
# isort: skip_file
|
||||
|
||||
# __torch_setup_begin__
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.data import DataLoader
|
||||
from torchvision import datasets
|
||||
from torchvision.transforms import ToTensor
|
||||
|
||||
def get_dataset():
|
||||
return datasets.FashionMNIST(
|
||||
root="/tmp/data",
|
||||
train=True,
|
||||
download=True,
|
||||
transform=ToTensor(),
|
||||
)
|
||||
|
||||
class NeuralNetwork(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.flatten = nn.Flatten()
|
||||
self.linear_relu_stack = nn.Sequential(
|
||||
nn.Linear(28 * 28, 512),
|
||||
nn.ReLU(),
|
||||
nn.Linear(512, 512),
|
||||
nn.ReLU(),
|
||||
nn.Linear(512, 10),
|
||||
)
|
||||
|
||||
def forward(self, inputs):
|
||||
inputs = self.flatten(inputs)
|
||||
logits = self.linear_relu_stack(inputs)
|
||||
return logits
|
||||
# __torch_setup_end__
|
||||
|
||||
# __torch_single_begin__
|
||||
def train_func():
|
||||
num_epochs = 3
|
||||
batch_size = 64
|
||||
|
||||
dataset = get_dataset()
|
||||
dataloader = DataLoader(dataset, batch_size=batch_size)
|
||||
|
||||
model = NeuralNetwork()
|
||||
|
||||
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()}")
|
||||
# __torch_single_end__
|
||||
|
||||
# __torch_distributed_begin__
|
||||
import ray.train.torch
|
||||
|
||||
def train_func_distributed():
|
||||
num_epochs = 3
|
||||
batch_size = 64
|
||||
|
||||
dataset = get_dataset()
|
||||
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
|
||||
dataloader = ray.train.torch.prepare_data_loader(dataloader)
|
||||
|
||||
model = NeuralNetwork()
|
||||
model = ray.train.torch.prepare_model(model)
|
||||
|
||||
criterion = nn.CrossEntropyLoss()
|
||||
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
|
||||
|
||||
for epoch in range(num_epochs):
|
||||
if ray.train.get_context().get_world_size() > 1:
|
||||
dataloader.sampler.set_epoch(epoch)
|
||||
|
||||
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()}")
|
||||
# __torch_distributed_end__
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# __torch_single_run_begin__
|
||||
train_func()
|
||||
# __torch_single_run_end__
|
||||
|
||||
# __torch_trainer_begin__
|
||||
from ray.train.torch import TorchTrainer
|
||||
from ray.train import ScalingConfig
|
||||
|
||||
# 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()
|
||||
# __torch_trainer_end__
|
||||
@@ -0,0 +1,159 @@
|
||||
import argparse
|
||||
import os
|
||||
import tempfile
|
||||
from typing import Tuple
|
||||
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
import ray
|
||||
import ray.train as train
|
||||
from ray.data import Dataset
|
||||
from ray.train import Checkpoint, DataConfig, ScalingConfig
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
|
||||
def get_datasets(split: float = 0.7) -> Tuple[Dataset]:
|
||||
dataset = ray.data.read_csv("s3://anonymous@air-example-data/regression.csv")
|
||||
|
||||
def combine_x(batch):
|
||||
return pd.DataFrame(
|
||||
{
|
||||
"x": batch[[f"x{i:03d}" for i in range(100)]].values.tolist(),
|
||||
"y": batch["y"],
|
||||
}
|
||||
)
|
||||
|
||||
dataset = dataset.map_batches(combine_x, batch_format="pandas")
|
||||
train_dataset, validation_dataset = dataset.repartition(
|
||||
num_blocks=4
|
||||
).train_test_split(split, shuffle=True)
|
||||
return train_dataset, validation_dataset
|
||||
|
||||
|
||||
def train_epoch(iterable_dataset, model, loss_fn, optimizer, device):
|
||||
model.train()
|
||||
for X, y in iterable_dataset:
|
||||
X = X.to(device)
|
||||
y = y.to(device)
|
||||
|
||||
# Compute prediction error
|
||||
pred = model(X)
|
||||
loss = loss_fn(pred, y)
|
||||
|
||||
# Backpropagation
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
|
||||
def validate_epoch(iterable_dataset, model, loss_fn, device):
|
||||
num_batches = 0
|
||||
model.eval()
|
||||
loss = 0
|
||||
with torch.no_grad():
|
||||
for X, y in iterable_dataset:
|
||||
X = X.to(device)
|
||||
y = y.to(device)
|
||||
num_batches += 1
|
||||
pred = model(X)
|
||||
loss += loss_fn(pred, y).item()
|
||||
loss /= num_batches
|
||||
result = {"loss": loss}
|
||||
return result
|
||||
|
||||
|
||||
def train_func(config):
|
||||
batch_size = config.get("batch_size", 32)
|
||||
hidden_size = config.get("hidden_size", 10)
|
||||
lr = config.get("lr", 1e-2)
|
||||
epochs = config.get("epochs", 3)
|
||||
|
||||
train_dataset_shard = train.get_dataset_shard("train")
|
||||
validation_dataset = train.get_dataset_shard("validation")
|
||||
|
||||
model = nn.Sequential(
|
||||
nn.Linear(100, hidden_size), nn.ReLU(), nn.Linear(hidden_size, 1)
|
||||
)
|
||||
model = train.torch.prepare_model(model)
|
||||
|
||||
loss_fn = nn.L1Loss()
|
||||
|
||||
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
|
||||
|
||||
results = []
|
||||
|
||||
def create_torch_iterator(shard):
|
||||
iterator = shard.iter_torch_batches(batch_size=batch_size)
|
||||
for batch in iterator:
|
||||
yield batch["x"].float(), batch["y"].float()
|
||||
|
||||
for _ in range(epochs):
|
||||
train_torch_dataset = create_torch_iterator(train_dataset_shard)
|
||||
validation_torch_dataset = create_torch_iterator(validation_dataset)
|
||||
|
||||
device = train.torch.get_device()
|
||||
|
||||
train_epoch(train_torch_dataset, model, loss_fn, optimizer, device)
|
||||
if train.get_context().get_world_rank() == 0:
|
||||
result = validate_epoch(validation_torch_dataset, model, loss_fn, device)
|
||||
else:
|
||||
result = {}
|
||||
results.append(result)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
torch.save(model.module.state_dict(), os.path.join(tmpdir, "model.pt"))
|
||||
train.report(result, checkpoint=Checkpoint.from_directory(tmpdir))
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def train_regression(num_workers=2, use_gpu=False):
|
||||
train_dataset, val_dataset = get_datasets()
|
||||
config = {"lr": 1e-2, "hidden_size": 20, "batch_size": 4, "epochs": 3}
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_loop_per_worker=train_func,
|
||||
train_loop_config=config,
|
||||
scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
|
||||
datasets={"train": train_dataset, "validation": val_dataset},
|
||||
dataset_config=DataConfig(datasets_to_split=["train"]),
|
||||
)
|
||||
|
||||
result = trainer.fit()
|
||||
return result
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--address", required=False, type=str, help="the address to use for Ray"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-workers",
|
||||
"-n",
|
||||
type=int,
|
||||
default=2,
|
||||
help="Sets number of workers for training.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--smoke-test",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Finish quickly for testing.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use-gpu", action="store_true", default=False, help="Use GPU for training."
|
||||
)
|
||||
|
||||
args, _ = parser.parse_known_args()
|
||||
|
||||
if args.smoke_test:
|
||||
# 2 workers, 1 for trainer, 1 for datasets
|
||||
ray.init(num_cpus=4)
|
||||
result = train_regression()
|
||||
else:
|
||||
ray.init(address=args.address)
|
||||
result = train_regression(num_workers=args.num_workers, use_gpu=args.use_gpu)
|
||||
print(result)
|
||||
Reference in New Issue
Block a user