TGraphX on Colab: Zero-Install Quickstart
Google Colab provides a free GPU runtime with Python, PyTorch, and most scientific Python libraries pre-installed. It is the fastest way to experiment with TGraphX without setting up a local environment. This guide covers everything needed to go from a blank Colab notebook to a running GNN experiment: installing TGraphX, building a graph, training a model, saving artifacts, and persisting results across sessions using Google Drive.
Installing TGraphX in Colab
Colab notebooks run in an ephemeral environment. Any package installed with !pip install lasts only for the current session. Start every session with the installation cell:
# Cell 1: Install TGraphX
!pip install tgraphx -q
# Verify installation
import tgraphx
print(f"TGraphX version: {tgraphx.__version__}")
import torch
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
The -q flag suppresses verbose output. After running this cell, all TGraphX modules are available for the session.
Checking the Runtime
Before starting, confirm that you have a GPU runtime. In Colab: Runtime > Change runtime type > GPU.
# Cell 2: Check hardware
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
if device.type == "cuda":
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
TGraphX works on CPU as well, so the quickstart runs on a CPU-only runtime. For large graphs or long experiments, the GPU runtime is worth enabling.
Building Your First Graph
# Cell 3: Create a TGraphX Graph object
import torch
from tgraphx import Graph
# A simple graph with 100 nodes and 300 edges
num_nodes = 100
num_edges = 300
node_features = torch.randn(num_nodes, 16) # 16-dimensional node features
edge_index = torch.randint(0, num_nodes, (2, num_edges), dtype=torch.long)
# Remove self-loops for clean experiments
mask = edge_index[0] != edge_index[1]
edge_index = edge_index[:, mask]
node_labels = torch.randint(0, 3, (num_nodes,)) # 3-class labels
g = Graph(
node_features=node_features,
edge_index=edge_index,
node_labels=node_labels,
)
print(f"Nodes: {g.node_features.shape[0]}")
print(f"Edges: {g.edge_index.shape[1]}")
print(f"Feature shape: {g.node_features.shape}")
Running the Shape Validator
TGraphX includes a tgraphx.doctor module that checks your graph for common configuration mistakes before training:
# Cell 4: Validate graph before training
from tgraphx.doctor import GraphDoctor
doctor = GraphDoctor(g)
report = doctor.diagnose()
for check, status in report.items():
icon = "OK" if status["passed"] else "WARN"
print(f"[{icon}] {check}: {status['message']}")
The doctor catches issues like isolated nodes, duplicate edges, inconsistent shapes, and features containing NaN values. Fixing these before training avoids silent failures during the forward pass. For more on shape validation, see the shape-aware validation guide.
Building and Training a GNN
# Cell 5: Define a simple two-layer GNN
import torch.nn as nn
import torch.nn.functional as F
from tgraphx.layers.sage import TensorGraphSAGELayer
class NodeClassifier(nn.Module):
def __init__(self, in_dim, hidden_dim, num_classes, dropout=0.5):
super().__init__()
self.conv1 = TensorGraphSAGELayer(in_dim, hidden_dim)
self.conv2 = TensorGraphSAGELayer(hidden_dim, num_classes)
self.dropout = dropout
def forward(self, x, edge_index):
x = F.relu(self.conv1(x, edge_index))
x = F.dropout(x, p=self.dropout, training=self.training)
return F.log_softmax(self.conv2(x, edge_index), dim=1)
model = NodeClassifier(in_dim=16, hidden_dim=32, num_classes=3).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
# Move graph to device
x = g.node_features.to(device)
edge_index = g.edge_index.to(device)
labels = g.node_labels.to(device)
# Simple 80/20 train/test split
num_nodes = x.shape[0]
perm = torch.randperm(num_nodes)
train_mask = torch.zeros(num_nodes, dtype=torch.bool, device=device)
test_mask = torch.zeros(num_nodes, dtype=torch.bool, device=device)
train_mask[perm[:80]] = True
test_mask[perm[80:]] = True
# Cell 6: Training loop
train_losses = []
test_accs = []
for epoch in range(100):
model.train()
optimizer.zero_grad()
log_probs = model(x, edge_index)
loss = F.nll_loss(log_probs[train_mask], labels[train_mask])
loss.backward()
optimizer.step()
train_losses.append(loss.item())
if epoch % 10 == 0:
model.eval()
with torch.no_grad():
log_probs = model(x, edge_index)
preds = log_probs[test_mask].argmax(dim=1)
acc = (preds == labels[test_mask]).float().mean().item()
test_accs.append(acc)
print(f"Epoch {epoch:3d} | Loss: {loss.item():.4f} | Test Acc: {acc:.4f}")
Plotting Results in Colab
Colab renders matplotlib plots inline:
# Cell 7: Plot training curves
import matplotlib.pyplot as plt
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
ax1.plot(train_losses)
ax1.set_title("Training Loss")
ax1.set_xlabel("Epoch")
ax1.set_ylabel("Loss")
ax1.grid(True, alpha=0.3)
ax2.plot(range(0, 100, 10), test_accs, marker='o')
ax2.set_title("Test Accuracy")
ax2.set_xlabel("Epoch")
ax2.set_ylabel("Accuracy")
ax2.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig("training_curves.png", dpi=150)
plt.show()
Saving Artifacts to Google Drive
Colab sessions are ephemeral. Anything saved to /content/ is lost when the session ends. Mount Google Drive for persistent storage:
# Cell 8: Mount Google Drive
from google.colab import drive
drive.mount("/content/drive")
# Create a directory for this project
import os
project_dir = "/content/drive/MyDrive/tgraphx_experiments"
os.makedirs(project_dir, exist_ok=True)
print(f"Project directory: {project_dir}")
Save model weights and results:
# Cell 9: Save model and results
import json
# Save model weights
model_path = f"{project_dir}/node_classifier.pt"
torch.save(model.state_dict(), model_path)
print(f"Model saved to: {model_path}")
# Save training history
history = {
"train_losses": train_losses,
"test_accs": test_accs,
"config": {
"in_dim": 16,
"hidden_dim": 32,
"num_classes": 3,
"lr": 0.01,
"epochs": 100,
}
}
with open(f"{project_dir}/history.json", "w") as f:
json.dump(history, f, indent=2)
# Save the training curves figure
plt.savefig(f"{project_dir}/training_curves.png", dpi=150, bbox_inches="tight")
print("Artifacts saved to Google Drive")
Common Colab-Specific Issues
Session disconnection. Colab disconnects after ~90 minutes of inactivity or ~12 hours total. Save checkpoints periodically inside your training loop:
if epoch % 25 == 0:
torch.save(model.state_dict(), f"{project_dir}/checkpoint_epoch{epoch}.pt")
Package version conflicts. Colab pre-installs specific versions of PyTorch and torchvision. TGraphX's PyPI release targets a minimum PyTorch version. If you see import errors after installation, restart the runtime (Runtime > Restart runtime) and re-run the install cell.
GPU memory exhaustion. The free Colab GPU has 12–16 GB of memory. For large tensor-valued node features ([N, C, H, W]), reduce batch size or use the chunk_size parameter in TensorGINLayer and TensorGraphSAGELayer to process edge batches rather than all edges at once.
Reinstalling every session. Because installations are ephemeral, pin the TGraphX version to avoid unexpected behavior when the package updates:
!pip install tgraphx==<version> -q # replace with your tested version
Random state is not preserved across cells. Each cell run is independent but shares Python state within a session. Set seeds at the top of your training cell rather than in a separate early cell, or use TGraphX's reproducibility module:
from tgraphx.reproducibility import ReproducibilityContext
ctx = ReproducibilityContext(seed=42)
ctx.seed_all()
Limitations and Honest Notes
Colab is not suitable for long training runs. The session time limit and potential disconnections make Colab unreliable for experiments exceeding a few hours. For multi-day training, use a dedicated server or cloud instance.
Google Drive mounting introduces latency. Saving and loading large files to Drive is slower than local disk. Cache large datasets to /content/ (ephemeral) during training and sync to Drive at checkpoints.
The free tier does not guarantee GPU availability. During periods of high demand, Colab may assign you a CPU runtime even when you select GPU. Check torch.cuda.is_available() at the start of each session.
For longer experiments with proper tracking, see the MLflow integration tutorial and the offline dashboard reports guide. For the full TGraphX package list, visit https://pypi.org/project/tgraphx/.
Scaling Up: From Quickstart to Real Research
The quickstart above uses a synthetic random graph. For real research in Colab, you will typically load a real dataset. Here is how to load a standard benchmark graph in TGraphX format:
# Cell 10: Load a real dataset (example: Cora via TGraphX datasets)
import torch
from tgraphx.datasets import load_planetoid
# Downloads Cora if not cached (cached to /content/ during session)
dataset = load_planetoid(name="Cora", root="/content/datasets/")
# Access the first (and only) graph
g = dataset[0]
print(f"Nodes: {g.node_features.shape[0]}") # 2708
print(f"Edges: {g.edge_index.shape[1]}") # 10556
print(f"Feature dim: {g.node_features.shape[1]}") # 1433
print(f"Classes: {g.node_labels.max().item() + 1}") # 7
For datasets too large to download in a Colab session, consider:
- Downloading to Google Drive in a separate Colab session, then mounting Drive and loading from there
- Using the TGraphX interop adapters to load from PyTorch Geometric datasets, which are more widely cached in Colab environments
Reproducibility in Colab
Colab introduces additional sources of non-reproducibility beyond standard PyTorch seeds: different sessions may have different CUDA versions, Colab updates Python library versions periodically, and GPU hardware assignments vary. To maximize reproducibility:
# Cell 11: Full reproducibility setup for Colab
from tgraphx.reproducibility import ReproducibilityContext
import subprocess
# Record environment
ctx = ReproducibilityContext(seed=42)
ctx.seed_all()
env = ctx.get_config()
print("Environment snapshot:")
for k, v in env.items():
print(f" {k}: {v}")
# Also capture the installed package versions
result = subprocess.run(["pip", "freeze"], capture_output=True, text=True)
with open(f"{project_dir}/requirements.txt", "w") as f:
f.write(result.stdout)
print("Requirements saved to Drive")
Saving the requirements.txt alongside your model weights means that someone reading your paper in two years can reconstruct the exact dependency versions you used. See the GNN research reproducibility guide for the full reproducibility workflow.
What This Article Builds On
This article is intentionally introductory — it assumes only familiarity with Python and Colab notebooks. For deeper coverage of the TGraphX concepts touched on here, see:
- Shape-aware validation guide — for understanding what GraphDoctor checks
- GNN research reproducibility guide — for full reproducibility setup
- The articles hub — for the full index of TGraphX tutorials
The Colab quickstart is a starting point, not an endpoint. Once your experiment grows beyond a few hundred epochs and a small synthetic graph, migrating to a local environment or a dedicated compute cluster will give you better control over hardware, session persistence, and debugging.
Frequently Asked Questions
Does TGraphX work on Colab's TPU runtime? No. TGraphX requires CUDA-capable GPU or CPU. TPU runtime uses a different compiler (XLA) that is not compatible with standard PyTorch operations out of the box. Use the GPU runtime.
Can I run TGraphX experiments in Colab Pro? Yes. Colab Pro and Pro+ provide longer session times and access to A100 GPUs, which is beneficial for experiments with large tensor-valued node features (e.g., [N, 3, 224, 224] image patches).
How do I share a Colab notebook with my TGraphX experiment? Click "Share" in the top-right corner. For reproducibility, make sure the notebook includes the installation cell and the reproducibility setup. A recipient who runs the notebook from top to bottom should get the same results, within the limits of library version differences.
Is there an official TGraphX Colab notebook template? Check the TGraphX repository at https://github.com/arashsajjadi/TGraphX for the latest Colab gallery. The colab_drafts/ directory contains starter notebooks for common tasks.
What if pip install tgraphx fails? This usually means a dependency conflict with a pre-installed library in Colab. Try pip install tgraphx --upgrade -q to force the latest compatible version, or pip install tgraphx --no-deps -q followed by installing missing dependencies individually. If problems persist, check the TGraphX GitHub issues page for known Colab compatibility notes specific to your current Colab runtime version.