Hi Everyone,
I am using TensorBoard in my PyTorch project and I faced a problem. After each training epoch I log histograms of model parameters using SummaryWriter.add_histogram(). When I set weight_decay=0 in my optimizer everything works as expected (see left image below). However, when I set weight decay to some non-zero value, e.g. weight_decay=0.1, the histogram writing time starts to explode with every epoch (right image), rendering training impossible. Along with that, all my RAM is being filled.
collage2800×1955 669 KB
Here is the code that demonstrates the problem:
import numpy as np
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.utils.data.dataloader import default_collate
from torch.utils.data.sampler import SubsetRandomSampler
from torchvision import datasets, transforms
from torch.utils.data import Subset
from torch.optim import Adam
from torch.optim.lr_scheduler import StepLR
from torch.utils.tensorboard import SummaryWriter

WEIGHT_DECAY = 0.1
EPOCHS = 10


class BaseDataLoader(DataLoader):
    """
    Base class for all data loaders
    """
    def __init__(self, dataset, batch_size, shuffle, validation_split, num_workers, collate_fn=default_collate):
        self.validation_split = validation_split
        self.shuffle = shuffle

        self.batch_idx = 0
        self.n_samples = len(dataset)

        self.sampler, self.valid_sampler = self._split_sampler(self.validation_split)

        self.init_kwargs = {
            'dataset': dataset,
            'batch_size': batch_size,
            'shuffle': self.shuffle,
            'collate_fn': collate_fn,
            'num_workers': num_workers
        }
        super().__init__(sampler=self.sampler, **self.init_kwargs)

    def _split_sampler(self, split):
        if split == 0.0:
            return None, None

        idx_full = np.arange(self.n_samples)

        np.random.seed(0)
        np.random.shuffle(idx_full)

        if isinstance(split, int):
            assert split > 0
            assert split < self.n_samples, "validation set size is configured to be larger than entire dataset."
            len_valid = split
        else:
            len_valid = int(self.n_samples * split)

        valid_idx = idx_full[0:len_valid]
        train_idx = np.delete(idx_full, np.arange(0, len_valid))

        train_sampler = SubsetRandomSampler(train_idx)
        valid_sampler = SubsetRandomSampler(valid_idx)

        # turn off shuffle option which is mutually exclusive with sampler
        self.shuffle = False
        self.n_samples = len(train_idx)

        return train_sampler, valid_sampler

    def split_validation(self):
        if self.valid_sampler is None:
            return None
        else:
            return DataLoader(sampler=self.valid_sampler, **self.init_kwargs)


class ADTargetTransform(object):
    """
    Transform target labels to 0 for normal, 1 for anomaly.
    """
    def __init__(self, normal_classes):
        self.normal_classes = normal_classes

    def __call__(self, target):
        return int(target not in self.normal_classes)


class MnistADDataLoader(BaseDataLoader):
    def __init__(self, data_dir, normal_classes, batch_size, oe_frac=0.0, shuffle=True, validation_split=0.0, num_workers=1, training=True):
        # Create normal and outlier classes tuples
        self.normal_classes = tuple(normal_classes)
        self.outlier_classes = tuple([x for x in list(range(0, 10)) if x not in self.normal_classes])

        trsfm = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])
        target_transform = ADTargetTransform(self.normal_classes)

        self.data_dir = data_dir
        mnist_dataset = datasets.MNIST(self.data_dir, train=training, download=True, transform=trsfm, target_transform=target_transform)

        # Get indexes of normal and outlier classes
        normal_idx = np.argwhere(np.isin(mnist_dataset.targets.cpu().data.numpy(), self.normal_classes)).flatten().tolist()
        outlier_idx = np.argwhere(np.isin(mnist_dataset.targets.cpu().data.numpy(), self.outlier_classes)).flatten().tolist()

        idx = normal_idx
        if oe_frac > 0:
            # Calculate how many outlier samples will be added to dataset
            n_outliers_to_use = int(round(len(normal_idx)*oe_frac/(1-oe_frac)))
            if n_outliers_to_use > len(outlier_idx):
                n_outliers_to_use = len(outlier_idx)
            idx += np.random.choice(outlier_idx, n_outliers_to_use).tolist()
            np.random.shuffle(idx)

        self.dataset = Subset(mnist_dataset, idx)
        super().__init__(self.dataset, batch_size, shuffle, validation_split, num_workers)


class MnistModel(nn.Module):
    def __init__(self, output_dim=10):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 150)
        self.fc2 = nn.Linear(150, 80)
        self.fc3 = nn.Linear(80, output_dim)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        out = self.fc3(x)
        return out


def loss_f(output, target):
    eps = 1e-9
    dists = torch.norm(output, p=2, dim=1) ** 2
    losses = torch.where(target == 0, dists, -torch.log(1 - torch.exp(-dists) + eps))
    return torch.mean(losses)


if __name__ == '__main__':
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    data_loader = MnistADDataLoader(data_dir="data/", normal_classes=[0], batch_size=128,
                                    oe_frac=0.1, shuffle=True, validation_split=0.1, num_workers=2)
    valid_data_loader = data_loader.split_validation()
    writer = SummaryWriter('log')
    model = MnistModel(10).to(device)
    trainable_params = filter(lambda p: p.requires_grad, model.parameters())
    optimizer = Adam(trainable_params, lr=0.001, weight_decay=WEIGHT_DECAY, amsgrad=True)
    lr_scheduler = StepLR(optimizer, step_size=50, gamma=0.1)
    for epoch in range(0, EPOCHS):
        print("Epoch", epoch, ":")
        model.train()
        for batch_idx, (data, target) in enumerate(data_loader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = loss_f(output, target)
            loss.backward()
            optimizer.step()
        start = time.time()
        for name, p in model.named_parameters():
            writer.add_histogram(name, p, bins='auto')
        end = time.time()
        print("\tHistograms:", end - start, "s")


I have killed quite some time to find the cause of my problem and a way to resolve it, however I still have no idea. I will be grateful if you can help me with that.
Thank you!