To train a PyTorch Transformer for multi-track music generation, the standard paradigm is to convert MIDI events into a sequential text-like language (tokens) that a decoder-only Transformer (like GPT) can read autoregressively.
To handle multiple tracks (e.g., Piano, Drums, Bass) in a single stream, we use an architecture inspired by frameworks like MMM (Multi-Track Music Machine) or MT3. Tracks are serialized sequentially using specific tokens to denote the track’s instrument type, followed by its notes:
[TRACK_START] [INSTRUMENT_PIANO] [NOTE_ON] [NOTE_OFF] … [TRACK_END] [TRACK_START] [INSTRUMENT_BASS] …
Below is a complete, self-contained Python code example implementing a dataset parser, a causal Transformer Decoder model, and an optimized training loop.
import os
import math
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import mido
# 1. TOKENIZER & VOCABULARY SETUP
# We define special tokens to serialize multiple tracks together.
SPECIAL_TOKENS = {
"<PAD>": 0,
"<BOS>": 1,
"<EOS>": 2,
"<TRACK_START>": 3,
"<TRACK_END>": 4,
}
class MultiTrackTokenizer:
def __init__(self):
self.token_to_id = dict(SPECIAL_TOKENS)
# Add Instrument tokens (MIDI program numbers 0-127)
for i in range(128):
self.token_to_id[f"INST_{i}"] = len(self.token_to_id)
# Add Note On / Note Off tokens (MIDI pitches 0-127)
for p in range(128):
self.token_to_id[f"NOTE_ON_{p}"] = len(self.token_to_id)
self.token_to_id[f"NOTE_OFF_{p}"] = len(self.token_to_id)
# Add Time Shift / Delta Time tokens (Quantized time steps, e.g., max 100 steps)
for t in range(100):
self.token_to_id[f"TIME_{t}"] = len(self.token_to_id)
self.id_to_token = {v: k for k, v in self.token_to_id.items()}
self.vocab_size = len(self.token_to_id)
def midi_to_tokens(self, midi_path, max_time_steps=100):
"""Converts a multi-track MIDI file into a flat sequence of token IDs."""
try:
mid = mido.MidiFile(midi_path)
except Exception as e:
return []
tokens = ["<BOS>"]
for track in mid.tracks:
track_tokens = []
current_instrument = 0 # Default to Piano
for msg in track:
# Catch program changes to update instrument tracking
if msg.type == 'program_change':
current_instrument = msg.program
# Encode delta time if present
if msg.time > 0:
quantized_time = min(int(msg.time / 10), max_time_steps - 1)
track_tokens.append(f"TIME_{quantized_time}")
# Encode notes
if msg.type == 'note_on' and msg.velocity > 0:
track_tokens.append(f"NOTE_ON_{msg.note}")
elif msg.type == 'note_off' or (msg.type == 'note_on' and msg.velocity == 0):
track_tokens.append(f"NOTE_OFF_{msg.note}")
# If the track contains musical notes, wrap it in multi-track tokens
if any("NOTE_" in t for t in track_tokens):
tokens.append("<TRACK_START>")
tokens.append(f"INST_{current_instrument}")
tokens.extend(track_tokens)
tokens.append("<TRACK_END>")
tokens.append("<EOS>")
return [self.token_to_id[t] for t in tokens if t in self.token_to_id]
# 2. PYTORCH DATASET
class MultiTrackMidiDataset(Dataset):
def __init__(self, midi_folder, tokenizer, seq_len=256):
self.tokenizer = tokenizer
self.seq_len = seq_len
self.samples = []
# Parse all MIDI files in directory and chunk them up
for file in os.listdir(midi_folder):
if file.endswith('.mid') or file.endswith('.midi'):
token_ids = self.tokenizer.midi_to_tokens(os.path.join(midi_folder, file))
# Chunk sequence into fixed lengths
for i in range(0, len(token_ids) - seq_len, seq_len // 2):
chunk = token_ids[i:i + seq_len + 1]
if len(chunk) == seq_len + 1:
self.samples.append(chunk)
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
# input (x) and target (y) shifted by one step for autoregressive training
tensor_data = torch.tensor(self.samples[idx], dtype=torch.long)
return tensor_data[:-1], tensor_data[1:]
# 3. TRANSFORMER MODEL ARCHITECTURE
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=512):
super().__init__()
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
self.register_buffer('pe', pe.unsqueeze(0))
def forward(self, x):
return x + self.pe[:, :x.size(1)]
class MusicTransformerLM(nn.Module):
def __init__(self, vocab_size, d_model=256, nhead=8, num_layers=4, dim_feedforward=512, max_seq_len=256):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
self.pos_encoder = PositionalEncoding(d_model, max_seq_len)
# Transformer Decoder structure
encoder_layer = nn.TransformerEncoderLayer(
d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward,
dropout=0.1, batch_first=True, activation='relu'
)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
self.fc_out = nn.Linear(d_model, vocab_size)
def generate_causal_mask(self, sz, device):
# Prevents positions from attending to subsequent positions (causal)
return torch.triu(torch.full((sz, sz), float('-inf'), device=device), diagonal=1)
def forward(self, src):
mask = self.generate_causal_mask(src.size(1), src.device)
x = self.embedding(src)
x = self.pos_encoder(x)
output = self.transformer(x, mask=mask, is_causal=True)
return self.fc_out(output)
# 4. TRAINING PIPELINE
def train_model(midi_dir, epochs=10, batch_size=16, seq_len=256):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
tokenizer = MultiTrackTokenizer()
# Dummy data generator setup if directory doesn't exist
if not os.path.exists(midi_dir) or len(os.listdir(midi_dir)) == 0:
print(f"Please create the directory '{midi_dir}' and add MIDI files.")
return
dataset = MultiTrackMidiDataset(midi_dir, tokenizer, seq_len=seq_len)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True)
model = MusicTransformerLM(vocab_size=tokenizer.vocab_size, max_seq_len=seq_len).to(device)
criterion = nn.CrossEntropyLoss(ignore_index=SPECIAL_TOKENS["<PAD>"])
optimizer = torch.optim.AdamW(model.parameters(), lr=0.0005)
model.train()
for epoch in range(epochs):
total_loss = 0
for batch_idx, (x, y) in enumerate(dataloader):
x, y = x.to(device), y.to(device)
optimizer.zero_grad()
outputs = model(x)
# Reshape output to (Batch * Seq, Vocab) and target to (Batch * Seq)
loss = criterion(outputs.view(-1, tokenizer.vocab_size), y.view(-1))
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(dataloader)
print(f"Epoch [{epoch+1}/{epochs}] | Loss: {avg_loss:.4f}")
# Save the trained model weights
torch.save(model.state_dict(), "multitrack_music_transformer.pt")
print("Training complete. Model saved!")
if __name__ == "__main__":
# Specify your MIDI dataset directory here
# Example: train_model("./my_midi_files", epochs=5)
pass
In the field of AI and Natural Language Processing (NLP), The Transformer is the engine…
If your goal is purely symbolic melody generation (e.g., generating MIDI files or notes rather…
According to PyTorch document, In PyTorch, torch.nn.functional.affine_grid generates a 2D or 3D flow field (sampling…
Creating an Android mobile app with a deep learning AI model involves several key steps:…
Click here to go to source files for R Machine Learning
Click here to download Python Machine Learning Source Files !