Wearable Activity Classification with SGLD¶

def class_distribution(dataset: ActivitySequenceDataset) -> Dict[str, int]:
    counts = torch.zeros(len(dataset.class_names), dtype=torch.int64)
    for label in dataset.labels:
        counts[label.item()] += 1
    return {dataset.class_names[i]: int(counts[i]) for i in range(len(counts))}

print('Train distribution:', class_distribution(train_dataset))
print('Test distribution :', class_distribution(test_dataset))

def plot_activity_examples(dataset: ActivitySequenceDataset, samples_per_class: int = 2) -> None:
    class_ids = list(dataset.class_names.keys())
    fig, axes = plt.subplots(samples_per_class, len(class_ids), figsize=(14, 4 * samples_per_class), sharex=True)
    if samples_per_class == 1:
        axes = axes[None, :]
    time_grid = dataset.time_grid.numpy()
    for col, class_id in enumerate(class_ids):
        candidates = (dataset.labels == class_id).nonzero(as_tuple=True)[0]
        picks = candidates[torch.randperm(len(candidates))[:samples_per_class]]
        for row, idx_tensor in enumerate(picks):
            idx = int(idx_tensor.item())
            series, label = dataset[idx]
            ax = axes[row, col]
            for axis_idx, axis_name in enumerate(['x', 'y', 'z']):
                ax.plot(time_grid, series[axis_idx].numpy(), label=f'axis {axis_name}')
            ax.set_title(f"class {dataset.class_names[label.item()]} (idx {idx})")
            ax.set_xlabel('time (normalised)')
            ax.set_ylabel('acceleration')
            ax.legend(loc='upper right')
    plt.tight_layout()

plot_activity_examples(train_dataset, samples_per_class=2)

class ResidualConvBlock(nn.Module):
    def __init__(self, channels: int, kernel_size: int = 5, dropout: float = 0.1) -> None:
        super().__init__()
        padding = kernel_size // 2
        self.block = nn.Sequential(
            nn.Conv1d(channels, channels, kernel_size=kernel_size, padding=padding),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Conv1d(channels, channels, kernel_size=kernel_size, padding=padding),
        )
        self.activation = nn.ReLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.activation(x + self.block(x))

class ActivityClassifier(nn.Module):
    def __init__(self, num_classes: int = 3, hidden_channels: int = 96, dropout: float = 0.2) -> None:
        super().__init__()
        self.stem = nn.Sequential(
            nn.Conv1d(3, hidden_channels, kernel_size=7, padding=3),
            nn.ReLU(),
            nn.Dropout(dropout),
        )
        self.encoder = nn.Sequential(
            ResidualConvBlock(hidden_channels, kernel_size=5, dropout=dropout),
            ResidualConvBlock(hidden_channels, kernel_size=5, dropout=dropout),
            nn.Conv1d(hidden_channels, hidden_channels, kernel_size=3, padding=1),
            nn.ReLU(),
        )
        self.pool = nn.AdaptiveAvgPool1d(1)
        self.head = nn.Linear(hidden_channels, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x)
        x = self.encoder(x)
        x = self.pool(x).squeeze(-1)
        return self.head(x)

model = ActivityClassifier().to(DEVICE)
total_params = sum(p.numel() for p in model.parameters())
print(model)
print(f'Total parameters: {total_params:,}')

model.eval()
mean_probs_list = []
var_probs_list = []
labels_list = []
inputs_list = []

with torch.inference_mode():
    for batch_inputs, batch_labels in test_loader:
        mean_probs, var_probs = predict_with_samples(
            model=model,
            samples=samples,
            x=batch_inputs,
            apply_softmax=True,
            device=DEVICE,
        )
        mean_probs_list.append(mean_probs)
        var_probs_list.append(var_probs)
        labels_list.append(batch_labels)
        inputs_list.append(batch_inputs)

mean_probs = torch.cat(mean_probs_list, dim=0)
var_probs = torch.cat(var_probs_list, dim=0)
true_labels = torch.cat(labels_list, dim=0)
inputs = torch.cat(inputs_list, dim=0)

pred_labels = mean_probs.argmax(dim=1)
test_accuracy = (pred_labels == true_labels).float().mean().item()
avg_variance = var_probs.mean(dim=1)

print(f'Posterior mean accuracy: {test_accuracy:.3f}')
print(f'Mean predictive variance: {avg_variance.mean().item():.4f}')

def compute_confusion_matrix(preds: torch.Tensor, targets: torch.Tensor, num_classes: int) -> torch.Tensor:
    matrix = torch.zeros(num_classes, num_classes, dtype=torch.int64)
    for true_label, pred_label in zip(targets, preds):
        matrix[true_label.item(), pred_label.item()] += 1
    return matrix

def plot_confusion_matrix(matrix: torch.Tensor, class_names) -> None:
    fig, ax = plt.subplots(figsize=(6, 5))
    im = ax.imshow(matrix.numpy(), cmap='Blues')
    ax.figure.colorbar(im, ax=ax)
    ax.set_xlabel('Predicted label')
    ax.set_ylabel('True label')
    ax.set_xticks(range(len(class_names)))
    ax.set_yticks(range(len(class_names)))
    ax.set_xticklabels(class_names)
    ax.set_yticklabels(class_names)
    for i in range(matrix.size(0)):
        for j in range(matrix.size(1)):
            ax.text(j, i, int(matrix[i, j].item()), ha="center", va="center", color="black")
    plt.tight_layout()

class_names = [train_dataset.class_names[i] for i in range(len(train_dataset.class_names))]
confusion = compute_confusion_matrix(pred_labels, true_labels, num_classes=len(class_names))
plot_confusion_matrix(confusion, class_names)

num_sequences = 6
half = num_sequences // 2
lowest_idx = torch.topk(avg_variance, k=half, largest=False).indices
highest_idx = torch.topk(avg_variance, k=half, largest=True).indices
selected = torch.cat([lowest_idx, highest_idx])
descriptors = ['low variance'] * half + ['high variance'] * half

time_axis = train_dataset.time_grid.numpy()
fig, axes = plt.subplots(2, half, figsize=(14, 6), sharex=True, sharey=True)
for ax, idx_tensor, desc in zip(axes.flatten(), selected, descriptors):
    idx = int(idx_tensor.item())
    sequence = inputs[idx].numpy()
    for channel, color in zip(range(sequence.shape[0]), ['tab:blue', 'tab:orange', 'tab:green']):
        ax.plot(time_axis, sequence[channel], color=color, alpha=0.8)
    title = (
        f"{desc} | true {train_dataset.class_names[true_labels[idx].item()]} | "
        f"p(run) {mean_probs[idx, 2].item():.2f} | var {avg_variance[idx].item():.4f}"
    )
    ax.set_title(title, fontsize=9)
    ax.set_xlabel('time')
    ax.set_ylabel('acceleration')
plt.tight_layout()