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
0 Comments