Two Moons

Example

import torch
import numpy as np
from sklearn.datasets import make_moons
from torch.utils.data import DataLoader, TensorDataset

# Generate samples
X, y = make_moons(n_samples=1000, noise=0.2, shuffle=True)
train_X, train_y = X[:800], y[:800]
test_X, test_y = X[800:], y[800:]

# Create the dataset
train_dataset = TensorDataset(torch.from_numpy(train_X).float(), torch.from_numpy(train_y).long())
test_dataset = TensorDataset(torch.from_numpy(test_X).float(), torch.from_numpy(test_y).long())

# Create the dataloaders
train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=True)

Visualizing

import matplotlib.pyplot as plt
from matplotlib import cm

# Plot the Two Moon dataset
plt.scatter(X[:, 0], X[:, 1], c=y)
plt.show()