分散トレーニング¶

from snowflake.ml.modeling.distributors.xgboost import XGBEstimator, XGBScalingConfig # Training parameters params = { 'objective': 'reg:squarederror', 'max_depth': 6, 'learning_rate': 0.1 } # Automatic scaling (recommended) estimator = XGBEstimator( params=params ) # Call with custom GPU scaling gpu_estimator = XGBEstimator( params=params, scaling_config=XGBScalingConfig(use_gpu=True) # optional - available resources will be used automatically ) # Train and evaluate booster = estimator.fit( dataset=train_connector, input_cols=['age', 'income', 'credit_score'], label_col='default_risk', eval_set=eval_connector, verbose_eval=10 ) # Access results booster = estimator.get_booster() # If you forgot to save the output of fit, get the booster from the estimator feature_importance = booster.get_score(importance_type='gain') 

# Create ShardedDataConnector for data ingestion from snowflake.ml.data.sharded_data_connector import ShardedDataConnector example_snowpark_dataframe = session.table("EXAMPLE_TRAINING_DATA") data_connector = ShardedDataConnector.from_dataframe(example_snowpark_dataframe) 

# Import necessary PyTorch libraries import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader # Define a simple neural network class SimpleNet(nn.Module): def __init__(self, input_size, hidden_size, output_size): super(SimpleNet, self).__init__() self.fc1 = nn.Linear(input_size, hidden_size) self.relu = nn.ReLU() self.fc2 = nn.Linear(hidden_size, output_size) def forward(self, x): x = self.fc1(x) x = self.relu(x) x = self.fc2(x) return x # Define the training function def train_func(): import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from snowflake.ml.modeling.distributors.pytorch import get_context # Use the Snowflake context to get the necessary methods to manage and retrieve information about the distributed training environment context = get_context() rank = context.get_rank() dist.init_process_group(backend='gloo') device = torch.device(f"cuda:{context.get_local_rank()}" if torch.cuda.is_available() else "cpu") # Initialize model, loss function, and optimizer model = SimpleNet(input_size=len(input_cols), hidden_size=32, output_size=1).to(device) model = DDP(model) criterion = nn.MSELoss() optimizer = optim.Adam(model.parameters(), lr=0.001) # Retrieve training data dataset_map = context.get_dataset_map() torch_dataset = dataset_map['train'].get_shard().to_torch_dataset(batch_size=1024) dataloader = DataLoader(torch_dataset) # Training loop for epoch in range(10): for batch_dict in dataloader: features = torch.cat([batch_dict[col].T for col in input_cols], dim=1).float().to(device) labels = batch_dict[label_col].T.squeeze(0).float().to(device) output = model(features) loss = criterion(output, labels.unsqueeze(1)) optimizer.zero_grad() loss.backward() optimizer.step() print(f'Epoch [{epoch+1}/10], Loss: {loss.item():.4f}') # Save the model to the model directory provided by the context if context.get_rank() == 0: torch.save( model.module.state_dict(), os.path.join(context.get_model_dir(), "model.pt") ) # Set up PyTorchDistributor for distributed training from snowflake.ml.modeling.distributors.pytorch import PyTorchDistributor, PyTorchScalingConfig, WorkerResourceConfig pytorch_trainer = PyTorchDistributor( train_func=train_func, # Optional Scaling Configuration, for single node multi-GPU training. scaling_config=PyTorchScalingConfig( num_nodes=1, num_workers_per_node=1, resource_requirements_per_worker=WorkerResourceConfig(num_cpus=0, num_gpus=4) ) ) # Run the training process pytorch_trainer.run(dataset_map={'train': data_connector}) 

response = pytorch_trainer.run( dataset_map={'train': data_connector}, artifact_stage_location="DB_NAME.SCHEMA_NAME.STAGE_NAME", ) stage_location = response.get_model_dir()