[ { "path": "LocallyConnected2d.py", "content": "\"\"\"\nTest implementation of locally connected 2d layer\nThe first part of the script was used for debugging\n\n@author: ptrblck\n\"\"\"\n\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn.modules.utils import _pair\n\n\n## DEBUG\nbatch_size = 5\nin_channels = 3\nh, w = 24, 24\nx = torch.ones(batch_size, in_channels, h, w)\nkh, kw = 3, 3 # kernel_size\ndh, dw = 1, 1 # stride\nx_windows = x.unfold(2, kh, dh).unfold(3, kw, dw)\nx_windows = x_windows.contiguous().view(*x_windows.size()[:-2], -1)\n\nout_channels = 2\nweights = torch.randn(1, out_channels, in_channels, *x_windows.size()[2:])\noutput = (x_windows.unsqueeze(1) * weights).sum([2, -1])\n## DEBUG\n\n\nclass LocallyConnected2d(nn.Module):\n def __init__(self, in_channels, out_channels, output_size, kernel_size, stride, bias=False):\n super(LocallyConnected2d, self).__init__()\n output_size = _pair(output_size)\n self.weight = nn.Parameter(\n torch.randn(1, out_channels, in_channels, output_size[0], output_size[1], kernel_size**2)\n )\n if bias:\n self.bias = nn.Parameter(\n torch.randn(1, out_channels, output_size[0], output_size[1])\n )\n else:\n self.register_parameter('bias', None)\n self.kernel_size = _pair(kernel_size)\n self.stride = _pair(stride)\n\n def forward(self, x):\n _, c, h, w = x.size()\n kh, kw = self.kernel_size\n dh, dw = self.stride\n x = x.unfold(2, kh, dh).unfold(3, kw, dw)\n x = x.contiguous().view(*x.size()[:-2], -1)\n # Sum in in_channel and kernel_size dims\n out = (x.unsqueeze(1) * self.weight).sum([2, -1])\n if self.bias is not None:\n out += self.bias\n return out\n\n\n# Create input\nbatch_size = 5\nin_channels = 3\nh, w = 24, 24\nx = torch.randn(batch_size, in_channels, h, w)\n\n# Create layer and test if backpropagation works\nout_channels = 2\noutput_size = 22\nkernel_size = 3\nstride = 1\nconv = LocallyConnected2d(\n in_channels, out_channels, output_size, kernel_size, stride, bias=True)\n\nout = conv(x)\nout.mean().backward()\nprint(conv.weight.grad)\n" }, { "path": "README.md", "content": "# PyTorch misc\nCollection of code snippets I've written for the [PyTorch discussion board](https://discuss.pytorch.org/).\n\nAll scripts were testes using the PyTorch 1.0 preview and torchvision `0.2.1`.\n\nAdditional libraries, e.g. `numpy` or `pandas`, are used in a few scripts.\n\nSome scripts might be a good starter to create a tutorial.\n\n## Overview\n\n * [accumulate_gradients](https://github.com/ptrblck/pytorch_misc/blob/master/accumulate_gradients.py) - Comparison of accumulated gradients/losses to vanilla batch update.\n * [adaptive_batchnorm](https://github.com/ptrblck/pytorch_misc/blob/master/adaptive_batchnorm.py)- Adaptive BN implementation using two additional parameters: `out = a * x + b * bn(x)`.\n * [adaptive_pooling_torchvision](https://github.com/ptrblck/pytorch_misc/blob/master/adaptive_pooling_torchvision.py) - Example of using adaptive pooling layers in pretrained models to use different spatial input shapes.\n * [batch_norm_manual](https://github.com/ptrblck/pytorch_misc/blob/master/batch_norm_manual.py) - Comparison of PyTorch BatchNorm layers and a manual calculation.\n * [change_crop_in_dataset](https://github.com/ptrblck/pytorch_misc/blob/master/change_crop_in_dataset.py) - Change the image crop size on the fly using a Dataset.\n * [channel_to_patches](https://github.com/ptrblck/pytorch_misc/blob/master/channel_to_patches.py) - Permute image data so that channel values of each pixel are flattened to an image patch around the pixel.\n * [conv_rnn](https://github.com/ptrblck/pytorch_misc/blob/master/conv_rnn.py) - Combines a 3DCNN with an RNN; uses windowed frames as inputs.\n * [csv_chunk_read](https://github.com/ptrblck/pytorch_misc/blob/master/csv_chunk_read.py) - Provide data chunks from continuous .csv file.\n * [densenet_forwardhook](https://github.com/ptrblck/pytorch_misc/blob/master/densenet_forwardhook.py) - Use forward hooks to get intermediate activations from `densenet121`. Uses separate modules to process these activations further.\n * [edge_weighting_segmentation](https://github.com/ptrblck/pytorch_misc/blob/master/edge_weighting_segmentation.py) - Apply weighting to edges for a segmentation task.\n * [image_rotation_with_matrix](https://github.com/ptrblck/pytorch_misc/blob/master/image_rotation_with_matrix.py) - Rotate an image given an angle using 1.) a nested loop and 2.) a rotation matrix and mesh grid.\n * [LocallyConnected2d](https://github.com/ptrblck/pytorch_misc/blob/master/LocallyConnected2d.py) - Implementation of a locally connected 2d layer.\n * [mnist_autoencoder](https://github.com/ptrblck/pytorch_misc/blob/master/mnist_autoencoder.py) - Simple autoencoder for MNIST data. Includes visualizations of output images, intermediate activations and conv kernels.\n * [mnist_permuted](https://github.com/ptrblck/pytorch_misc/blob/master/mnist_permuted.py) - MNIST training using permuted pixel locations.\n * [model_sharding_data_parallel](https://github.com/ptrblck/pytorch_misc/blob/master/model_sharding_data_parallel.py) - Model sharding with `DataParallel` using 2 pairs of 2 GPUs.\n * [momentum_update_nograd](https://github.com/ptrblck/pytorch_misc/blob/master/momentum_update_nograd.py) - Script to see how parameters are updated when an optimizer is used with momentum/running estimates, even if gradients are zero.\n * [pytorch_redis](https://github.com/ptrblck/pytorch_misc/blob/master/pytorch_redis.py) - Script to demonstrate the loading data from redis using a PyTorch Dataset and DataLoader.\n * [shared_array](https://github.com/ptrblck/pytorch_misc/blob/master/shared_array.py) - Script to demonstrate the usage of shared arrays using multiple workers.\n * [shared_dict](https://github.com/ptrblck/pytorch_misc/blob/master/shared_dict.py) - Script to demonstrate the usage of shared dicts using multiple workers.\n * [unet_demo](https://github.com/ptrblck/pytorch_misc/blob/master/unet_demo.py) - Simple UNet demo.\n * [weighted_sampling](https://github.com/ptrblck/pytorch_misc/blob/master/weighted_sampling.py) - Usage of WeightedRandomSampler using an imbalanced dataset with class imbalance 99 to 1.\n\n\nFeedback is very welcome!\n" }, { "path": "accumulate_gradients.py", "content": "\"\"\"\nComparison of accumulated gradients/losses to vanilla batch update.\nComments from @albanD:\nhttps://discuss.pytorch.org/t/why-do-we-need-to-set-the-gradients-manually-to-zero-in-pytorch/4903/20\n\n@author: ptrblck\n\"\"\"\n\nimport torch\nimport torch.nn as nn\n\n\n# Accumulate loss for each samples\n# more runtime, more memory\nx1 = torch.ones(2, 1)\nw1 = torch.ones(1, 1, requires_grad=True)\ny1 = torch.ones(2, 1) * 2\n\ncriterion = nn.MSELoss()\n\nloss1 = 0\nfor i in range(10):\n output1 = torch.matmul(x1, w1)\n loss1 += criterion(output1, y1)\nloss1 /= 10 # scale loss to match batch gradient\nloss1.backward()\n\nprint('Accumulated losses: {}'.format(w1.grad))\n\n# Use whole batch to calculate gradient\n# least runtime, more memory\nx2 = torch.ones(20, 1)\nw2 = torch.ones(1, 1, requires_grad=True)\ny2 = torch.ones(20, 1) * 2\n\noutput2 = torch.matmul(x2, w2)\nloss2 = criterion(output2, y2)\nloss2.backward()\nprint('Batch gradient: {}'.format(w2.grad))\n\n# Accumulate scaled gradient\n# more runtime, least memory\nx3 = torch.ones(2, 1)\nw3 = torch.ones(1, 1, requires_grad=True)\ny3 = torch.ones(2, 1) * 2\n\nfor i in range(10):\n output3 = torch.matmul(x3, w3)\n loss3 = criterion(output3, y3)\n loss3 /= 10\n loss3.backward()\n\nprint('Accumulated gradient: {}'.format(w3.grad))\n" }, { "path": "adaptive_batchnorm.py", "content": "\"\"\"\nImplementation of Adaptive BatchNorm\n\n@author: ptrblck\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torchvision import datasets, transforms\n\n\n# Globals\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'\nseed = 2809\nbatch_size = 10\nlr = 0.01\nlog_interval = 10\nepochs = 10\ntorch.manual_seed(seed)\n\n\nclass AdaptiveBatchNorm2d(nn.Module):\n '''\n Adaptive BN implementation using two additional parameters:\n out = a * x + b * bn(x)\n '''\n def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True):\n super(AdaptiveBatchNorm2d, self).__init__()\n self.bn = nn.BatchNorm2d(num_features, eps, momentum, affine)\n self.a = nn.Parameter(torch.FloatTensor(1, 1, 1, 1))\n self.b = nn.Parameter(torch.FloatTensor(1, 1, 1, 1))\n\n def forward(self, x):\n return self.a * x + self.b * self.bn(x)\n\n\nclass MyNet(nn.Module):\n def __init__(self):\n super(MyNet, self).__init__()\n self.conv1 = nn.Conv2d(in_channels=1,\n out_channels=10,\n kernel_size=5)\n self.conv1_bn = AdaptiveBatchNorm2d(10)\n self.conv2 = nn.Conv2d(in_channels=10,\n out_channels=20,\n kernel_size=5)\n self.conv2_bn = AdaptiveBatchNorm2d(20)\n self.fc1 = nn.Linear(320, 50)\n self.fc2 = nn.Linear(50, 10)\n\n def forward(self, x):\n x = F.relu(F.max_pool2d(self.conv1_bn(self.conv1(x)), 2))\n x = F.relu(F.max_pool2d(self.conv2_bn(self.conv2(x)), 2))\n x = x.view(-1, 320)\n x = F.relu(self.fc1(x))\n x = self.fc2(x)\n return F.log_softmax(x, dim=1)\n\n\ndef train(epoch):\n model.train()\n for batch_idx, (data, target) in enumerate(train_loader):\n data, target = data.to(device), target.to(device)\n optimizer.zero_grad()\n output = model(data)\n loss = F.nll_loss(output, target)\n loss.backward()\n optimizer.step()\n if batch_idx % log_interval == 0:\n print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n epoch, batch_idx * len(data), len(train_loader.dataset),\n 100. * batch_idx / len(train_loader), loss.item()))\n\n\ndef test():\n model.eval()\n test_loss = 0\n correct = 0\n with torch.no_grad():\n for data, target in test_loader:\n data, target = data.to(device), target.to(device)\n output = model(data)\n # sum up batch loss\n test_loss += F.nll_loss(output, target, size_average=False).item()\n # get the index of the max log-probability\n pred = output.data.max(1, keepdim=True)[1]\n correct += pred.eq(target.data.view_as(pred)).cpu().sum()\n\n test_loss /= len(test_loader.dataset)\n print('\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n'.format(\n test_loss, correct, len(test_loader.dataset),\n 100. * correct / len(test_loader.dataset)))\n\n\ntrain_loader = torch.utils.data.DataLoader(\n datasets.MNIST('./data', train=True, download=True,\n transform=transforms.Compose([\n transforms.ToTensor(),\n transforms.Normalize((0.1307,), (0.3081,))\n ])),\n batch_size=batch_size,\n shuffle=True)\n\ntest_loader = torch.utils.data.DataLoader(\n datasets.MNIST('./data', train=False, transform=transforms.Compose([\n transforms.ToTensor(),\n transforms.Normalize((0.1307,), (0.3081,))\n ])),\n batch_size=batch_size,\n shuffle=True)\n\nmodel = MyNet().to(device)\noptimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.5)\n\nfor epoch in range(1, epochs + 1):\n train(epoch)\n test()\n" }, { "path": "adaptive_pooling_torchvision.py", "content": "\"\"\"\nAdaptive pooling layer examples\n\n@author: ptrblck\n\"\"\"\n\nimport torch\nimport torch.nn as nn\n\nimport torchvision.models as models\n\n\n# Use standard model with [batch_size, 3, 224, 224] input\nmodel = models.vgg16(pretrained=False)\nbatch_size = 1\nx = torch.randn(batch_size, 3, 224, 224)\noutput = model(x)\n\n# Try bigger input\nx_big = torch.randn(batch_size, 3, 299, 299)\ntry:\n output = model(x_big)\nexcept RuntimeError as e:\n print(e)\n\n# Try smaller input\nx_small = torch.randn(batch_size, 3, 128, 128)\ntry:\n output = model(x_small)\nexcept RuntimeError as e:\n print(e)\n# Both don't work, since we get a size mismatch for these sizes\n\n# Get the size of the last activation map before the classifier\ndef size_hook(module, input, output):\n print(output.shape)\n\nmodel.features[-1].register_forward_hook(size_hook)\noutput = model(x)\n\n# We see that the last pooling layer returns an activation of\n# [batch_size, 512, 7, 7]. So let's replace it with an adaptive layer with an\n# output shape of 7x7.\nmodel.features[-1] = nn.AdaptiveMaxPool2d(output_size=7)\n\n# Now let's try the other shapes again\noutput = model(x_big)\noutput = model(x_small)\n\nx_tiny = torch.randn(batch_size, 3, 16, 16)\noutput = model(x_tiny)\n\n# Now these inputs are working!\n# There is however a minimal size as we need a spatial size of at least 1x1\n# to pass into the adaptive pooling layer\nx_too_small = torch.randn(batch_size, 3, 15, 15)\ntry:\n output = model(x_too_small)\nexcept RuntimeError as e:\n print(e)\n" }, { "path": "batch_norm_manual.py", "content": "\"\"\"\nComparison of manual BatchNorm2d layer implementation in Python and\nnn.BatchNorm2d\n\n@author: ptrblck\n\"\"\"\n\nimport torch\nimport torch.nn as nn\n\n\ndef compare_bn(bn1, bn2):\n err = False\n if not torch.allclose(bn1.running_mean, bn2.running_mean):\n print('Diff in running_mean: {} vs {}'.format(\n bn1.running_mean, bn2.running_mean))\n err = True\n\n if not torch.allclose(bn1.running_var, bn2.running_var):\n print('Diff in running_var: {} vs {}'.format(\n bn1.running_var, bn2.running_var))\n err = True\n\n if bn1.affine and bn2.affine:\n if not torch.allclose(bn1.weight, bn2.weight):\n print('Diff in weight: {} vs {}'.format(\n bn1.weight, bn2.weight))\n err = True\n\n if not torch.allclose(bn1.bias, bn2.bias):\n print('Diff in bias: {} vs {}'.format(\n bn1.bias, bn2.bias))\n err = True\n\n if not err:\n print('All parameters are equal!')\n\n\nclass MyBatchNorm2d(nn.BatchNorm2d):\n def __init__(self, num_features, eps=1e-5, momentum=0.1,\n affine=True, track_running_stats=True):\n super(MyBatchNorm2d, self).__init__(\n num_features, eps, momentum, affine, track_running_stats)\n\n def forward(self, input):\n self._check_input_dim(input)\n\n exponential_average_factor = 0.0\n\n if self.training and self.track_running_stats:\n if self.num_batches_tracked is not None:\n self.num_batches_tracked += 1\n if self.momentum is None: # use cumulative moving average\n exponential_average_factor = 1.0 / float(self.num_batches_tracked)\n else: # use exponential moving average\n exponential_average_factor = self.momentum\n\n # calculate running estimates\n if self.training:\n mean = input.mean([0, 2, 3])\n # use biased var in train\n var = input.var([0, 2, 3], unbiased=False)\n n = input.numel() / input.size(1)\n with torch.no_grad():\n self.running_mean = exponential_average_factor * mean\\\n + (1 - exponential_average_factor) * self.running_mean\n # update running_var with unbiased var\n self.running_var = exponential_average_factor * var * n / (n - 1)\\\n + (1 - exponential_average_factor) * self.running_var\n else:\n mean = self.running_mean\n var = self.running_var\n\n input = (input - mean[None, :, None, None]) / (torch.sqrt(var[None, :, None, None] + self.eps))\n if self.affine:\n input = input * self.weight[None, :, None, None] + self.bias[None, :, None, None]\n\n return input\n\n\n# Init BatchNorm layers\nmy_bn = MyBatchNorm2d(3, affine=True)\nbn = nn.BatchNorm2d(3, affine=True)\n\ncompare_bn(my_bn, bn) # weight and bias should be different\n# Load weight and bias\nmy_bn.load_state_dict(bn.state_dict())\ncompare_bn(my_bn, bn)\n\n# Run train\nfor _ in range(10):\n scale = torch.randint(1, 10, (1,)).float()\n bias = torch.randint(-10, 10, (1,)).float()\n x = torch.randn(10, 3, 100, 100) * scale + bias\n out1 = my_bn(x)\n out2 = bn(x)\n compare_bn(my_bn, bn)\n\n torch.allclose(out1, out2)\n print('Max diff: ', (out1 - out2).abs().max())\n\n# Run eval\nmy_bn.eval()\nbn.eval()\nfor _ in range(10):\n scale = torch.randint(1, 10, (1,)).float()\n bias = torch.randint(-10, 10, (1,)).float()\n x = torch.randn(10, 3, 100, 100) * scale + bias\n out1 = my_bn(x)\n out2 = bn(x)\n compare_bn(my_bn, bn)\n\n torch.allclose(out1, out2)\n print('Max diff: ', (out1 - out2).abs().max())\n" }, { "path": "change_crop_in_dataset.py", "content": "\"\"\"\nChange the crop size on the fly using a Dataset.\nMyDataset.set_state(stage) switches between crop sizes.\nAlternatively, the crop size could be specified.\n\n@author: ptrblck\n\"\"\"\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nimport torchvision.transforms.functional as TF\n\n\nclass MyDataset(Dataset):\n def __init__(self):\n self.images = [TF.to_pil_image(x) for x in torch.ByteTensor(10, 3, 48, 48)]\n self.set_stage(0)\n\n def __getitem__(self, index):\n image = self.images[index]\n\n # Switch your behavior depending on stage\n image = self.crop(image)\n x = TF.to_tensor(image)\n return x\n\n def set_stage(self, stage):\n if stage == 0:\n print('Using (32, 32) crops')\n self.crop = transforms.RandomCrop((32, 32))\n elif stage == 1:\n print('Using (28, 28) crops')\n self.crop = transforms.RandomCrop((28, 28))\n\n def __len__(self):\n return len(self.images)\n\n\ndataset = MyDataset()\nloader = DataLoader(dataset,\n batch_size=2,\n num_workers=2,\n shuffle=True)\n\n# Use standard crop size\nfor batch_idx, data in enumerate(loader):\n print('Batch idx {}, data shape {}'.format(\n batch_idx, data.shape))\n\n# Switch to stage1 crop size\nloader.dataset.set_stage(1)\n\n# Check the shape again\nfor batch_idx, data in enumerate(loader):\n print('Batch idx {}, data shape {}'.format(\n batch_idx, data.shape))\n" }, { "path": "channel_to_patches.py", "content": "\"\"\"\nPermute image data so that channel values of each pixel are flattened to an image patch around the pixel.\n\n@author: ptrblck\n\"\"\"\n\nimport torch\n\nB, C, H, W = 2, 16, 4, 4\n# Create dummy input with same values in each channel\nx = torch.arange(C)[None, :, None, None].repeat(B, 1, H, W)\nprint(x)\n# Permute channel dimension to last position and view as 4x4 windows\nx = x.permute(0, 2, 3, 1).view(B, H, W, 4, 4)\nprint(x)\n# Permute \"window dims\" with spatial dims, view as desired output\nx = x.permute(0, 1, 3, 2, 4).contiguous().view(B, 1, 4*H, 4*W)\nprint(x)\n" }, { "path": "conv_rnn.py", "content": "\"\"\"\nCombine Conv3d with an RNN Module.\nUse windowed frames as inputs.\n\n@author: ptrblck\n\"\"\"\n\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset\n\n\nclass MyModel(nn.Module):\n def __init__(self, window=16):\n super(MyModel, self).__init__()\n self.conv_model = nn.Sequential(\n nn.Conv3d(\n in_channels=3,\n out_channels=6,\n kernel_size=3,\n stride=1,\n padding=1),\n nn.MaxPool3d((1, 2, 2)),\n nn.ReLU()\n )\n\n self.rnn = nn.RNN(\n input_size=6*16*12*12,\n hidden_size=1,\n num_layers=1,\n batch_first=True\n )\n self.hidden = torch.zeros(1, 1, 1)\n self.window = window\n\n def forward(self, x):\n self.hidden = torch.zeros(1, 1, 1) # reset hidden\n\n activations = []\n for idx in range(0, x.size(2), self.window):\n x_ = x[:, :, idx:idx+self.window]\n x_ = self.conv_model(x_)\n x_ = x_.view(x_.size(0), 1, -1)\n activations.append(x_)\n x = torch.cat(activations, 1)\n out, hidden = self.rnn(x, self.hidden)\n\n return out, hidden\n\n\nclass MyDataset(Dataset):\n '''\n Returns windowed frames from sequential data.\n '''\n def __init__(self, frames=512):\n self.data = torch.randn(3, 2048, 24, 24)\n self.frames = frames\n\n def __getitem__(self, index):\n index = index * self.frames\n x = self.data[:, index:index+self.frames]\n return x\n\n def __len__(self):\n return self.data.size(1) / self.frames\n\n\nmodel = MyModel()\ndataset = MyDataset()\nx = dataset[0]\noutput, hidden = model(x.unsqueeze(0))\n" }, { "path": "csv_chunk_read.py", "content": "\"\"\"\nProvide data chunks from continuous .csv file.\n\n@author: ptrblck\n\"\"\"\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\n\nimport numpy as np\nimport pandas as pd\n\n# Create dummy csv data\nnb_samples = 110\na = np.arange(nb_samples)\ndf = pd.DataFrame(a, columns=['data'])\ndf.to_csv('data.csv', index=False)\n\n\n# Create Dataset\nclass CSVDataset(Dataset):\n def __init__(self, path, chunksize, nb_samples):\n self.path = path\n self.chunksize = chunksize\n self.len = nb_samples // self.chunksize\n\n def __getitem__(self, index):\n '''\n Get next chunk of data\n '''\n x = next(\n pd.read_csv(\n self.path,\n skiprows=index * self.chunksize + 1, # +1, since we skip the header\n chunksize=self.chunksize,\n names=['data']))\n x = torch.from_numpy(x.data.values)\n return x\n\n def __len__(self):\n return self.len\n\n\ndataset = CSVDataset('data.csv', chunksize=10, nb_samples=nb_samples)\nloader = DataLoader(dataset, batch_size=10, num_workers=1, shuffle=False)\n\nfor batch_idx, data in enumerate(loader):\n print('batch: {}\\tdata: {}'.format(batch_idx, data))\n" }, { "path": "densenet_forwardhook.py", "content": "\"\"\"\nUse forward hooks to get intermediate activations from densenet121.\nCreate additional conv layers to process these activations to get a\ndesired number of output channels\n\n@author: ptrblck\n\"\"\"\n\nimport torch\nimport torch.nn as nn\n\nfrom torchvision import models\n\n\nactivations = {}\ndef get_activation(name):\n def hook(model, input, output):\n activations[name] = output\n return hook\n\n# Create Model\nmodel = models.densenet121(pretrained=False)\n\n# Register forward hooks with name\nfor name, child in model.features.named_children():\n if 'denseblock' in name:\n print(name)\n child.register_forward_hook(get_activation(name))\n\n# Forward pass\nx = torch.randn(1, 3, 224, 224)\noutput = model(x)\n\n# Create convs to get desired out_channels\nout_channels = 1\nconvs = {'denseblock1': nn.Conv2d(256, out_channels, 1,),\n 'denseblock2': nn.Conv2d(512, out_channels, 1),\n 'denseblock3': nn.Conv2d(1024, out_channels, 1),\n 'denseblock4': nn.Conv2d(1024, out_channels, 1)}\n\n# Apply conv on each activation\nfor key in activations:\n act = activations[key]\n act = convs[key](act)\n print(key, act.shape)\n" }, { "path": "edge_weighting_segmentation.py", "content": "\"\"\"\nApply weighting to edges for a segmentation task\n\n@author: ptrblck\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nimport matplotlib.pyplot as plt\n\n\n# Create dummy input and target with two squares\noutput = F.log_softmax(torch.randn(1, 3, 24, 24), 1)\ntarget = torch.zeros(1, 24, 24, dtype=torch.long)\ntarget[0, 4:12, 4:12] = 1\ntarget[0, 14:20, 14:20] = 2\nplt.imshow(target[0])\n\n# Edge calculation\n# Get binary target\nbin_target = torch.where(target > 0, torch.tensor(1), torch.tensor(0))\nplt.imshow(bin_target[0])\n\n# Use average pooling to get edge\no = F.avg_pool2d(bin_target.float(), kernel_size=3, padding=1, stride=1)\nplt.imshow(o[0])\n\nedge_idx = (o.ge(0.01) * o.le(0.99)).float()\nplt.imshow(edge_idx[0])\n\n# Create weight mask\nweights = torch.ones_like(edge_idx, dtype=torch.float32)\nweights_sum0 = weights.sum() # Save initial sum for later rescaling\nweights = weights + edge_idx * 2. # Weight edged with 2x loss\nweights_sum1 = weights.sum()\nweights = weights / weights_sum1 * weights_sum0 # Rescale weigths\nplt.imshow(weights[0])\n\n# Calculate loss\ncriterion = nn.NLLLoss(reduction='none')\nloss = criterion(output, target)\nloss = loss * weights # Apply weighting\nloss = loss.sum() / weights.sum() # Scale loss\n" }, { "path": "image_rotation_with_matrix.py", "content": "\"\"\"\nRotate image given an angle.\n1. Calculate rotated position for each input pixel\n2. Use meshgrid and rotation matrix to achieve the same\n\n@author: ptrblck\n\"\"\"\n\nimport torch\nimport numpy as np\n\n\n# Create dummy image\nbatch_size = 1\nim = torch.zeros(batch_size, 1, 10, 10)\nim[:, :, :, 2] = 1.\n\n# Set angle\nangle = torch.tensor([72 * np.pi / 180.])\n\n# Calculate rotation for each target pixel\nx_mid = (im.size(2) + 1) / 2.\ny_mid = (im.size(3) + 1) / 2.\nim_rot = torch.zeros_like(im)\nfor r in range(im.size(2)):\n for c in range(im.size(3)):\n x = (r - x_mid) * torch.cos(angle) + (c - y_mid) * torch.sin(angle)\n y = -1.0 * (r - x_mid) * torch.sin(angle) + (c - y_mid) * torch.cos(angle)\n x = torch.round(x) + x_mid\n y = torch.round(y) + y_mid\n\n if (x >= 0 and y >= 0 and x < im.size(2) and y < im.size(3)):\n im_rot[:, :, r, c] = im[:, :, x.long().item(), y.long().item()]\n\n\n# Calculate rotation with inverse rotation matrix\nrot_matrix = torch.tensor([[torch.cos(angle), torch.sin(angle)],\n [-1.0*torch.sin(angle), torch.cos(angle)]])\n\n# Use meshgrid for pixel coords\nxv, yv = torch.meshgrid(torch.arange(im.size(2)), torch.arange(im.size(3)))\nxv = xv.contiguous()\nyv = yv.contiguous()\nsrc_ind = torch.cat((\n (xv.float() - x_mid).view(-1, 1),\n (yv.float() - y_mid).view(-1, 1)),\n dim=1\n)\n\n# Calculate indices using rotation matrix\nsrc_ind = torch.matmul(src_ind, rot_matrix.t())\nsrc_ind = torch.round(src_ind)\nsrc_ind += torch.tensor([[x_mid, y_mid]])\n\n# Set out of bounds indices to limits\nsrc_ind[src_ind < 0] = 0.\nsrc_ind[:, 0][src_ind[:, 0] >= im.size(2)] = float(im.size(2)) - 1\nsrc_ind[:, 1][src_ind[:, 1] >= im.size(3)] = float(im.size(3)) - 1\n\n# Create new rotated image\nim_rot2 = torch.zeros_like(im)\nsrc_ind = src_ind.long()\nim_rot2[:, :, xv.view(-1), yv.view(-1)] = im[:, :, src_ind[:, 0], src_ind[:, 1]]\nim_rot2 = im_rot2.view(batch_size, 1, 10, 10)\n\nprint('Using method 1: {}'.format(im_rot))\nprint('Using method 2: {}'.format(im_rot2))\n" }, { "path": "mnist_autoencoder.py", "content": "\"\"\"\nSimple autoencoder for MNIST data.\nVisualizes some output images, intermediate activations as well as some conv\nkernels.\n\n@author: ptrblck\n\"\"\"\n\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader\n\nimport torchvision.transforms as transforms\nimport torchvision.datasets as datasets\n\nimport matplotlib.pyplot as plt\n\n\nclass MyModel(nn.Module):\n def __init__(self):\n super(MyModel, self).__init__()\n self.conv1 = nn.Conv2d(1, 3, 3, 1, 1)\n self.pool1 = nn.MaxPool2d(2)\n self.conv2 = nn.Conv2d(3, 6, 3, 1, 1)\n self.pool2 = nn.MaxPool2d(2)\n\n self.conv_trans1 = nn.ConvTranspose2d(6, 3, 4, 2, 1)\n self.conv_trans2 = nn.ConvTranspose2d(3, 1, 4, 2, 1)\n\n def forward(self, x):\n x = F.relu(self.pool1(self.conv1(x)))\n x = F.relu(self.pool2(self.conv2(x)))\n x = F.relu(self.conv_trans1(x))\n x = self.conv_trans2(x)\n return x\n\ndataset = datasets.MNIST(\n root='./data',\n transform=transforms.ToTensor()\n)\nloader = DataLoader(\n dataset,\n num_workers=2,\n batch_size=8,\n shuffle=True\n)\n\nmodel = MyModel()\ncriterion = nn.BCEWithLogitsLoss()\noptimizer = optim.Adam(model.parameters(), lr=1e-3)\n\nepochs = 1\nfor epoch in range(epochs):\n for batch_idx, (data, target) in enumerate(loader):\n optimizer.zero_grad()\n output = model(data)\n loss = criterion(output, data)\n loss.backward()\n optimizer.step()\n\n print('Epoch {}, Batch idx {}, loss {}'.format(\n epoch, batch_idx, loss.item()))\n\n\ndef normalize_output(img):\n img = img - img.min()\n img = img / img.max()\n return img\n\n# Plot some images\nidx = torch.randint(0, output.size(0), ())\npred = normalize_output(output[idx, 0])\nimg = data[idx, 0]\n\nfig, axarr = plt.subplots(1, 2)\naxarr[0].imshow(img.detach().numpy())\naxarr[1].imshow(pred.detach().numpy())\n\n# Visualize feature maps\nactivation = {}\ndef get_activation(name):\n def hook(model, input, output):\n activation[name] = output.detach()\n return hook\n\nmodel.conv1.register_forward_hook(get_activation('conv1'))\ndata, _ = dataset[0]\ndata.unsqueeze_(0)\noutput = model(data)\n\nact = activation['conv1'].squeeze()\nfig, axarr = plt.subplots(act.size(0))\nfor idx in range(act.size(0)):\n axarr[idx].imshow(act[idx])\n\n# Visualize conv filter\nkernels = model.conv1.weight.detach()\nfig, axarr = plt.subplots(kernels.size(0))\nfor idx in range(kernels.size(0)):\n axarr[idx].imshow(kernels[idx].squeeze())\n" }, { "path": "mnist_permuted.py", "content": "\"\"\"\nPermute all pixels of MNIST data and try to learn it using simple model.\n\n@author: ptrblck\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader\nimport torch.nn.functional as F\n\nfrom torchvision import datasets\nfrom torchvision import transforms\n\nimport numpy as np\n\n# Create random indices to permute images\nindices = np.arange(28*28)\nnp.random.shuffle(indices)\n\n\ndef shuffle_image(tensor):\n tensor = tensor.view(-1)[indices].view(1, 28, 28)\n return tensor\n\n\n# Apply permuatation using transforms.Lambda\ntrain_dataset = datasets.MNIST(root='./data',\n download=False,\n train=True,\n transform=transforms.Compose([\n transforms.ToTensor(),\n transforms.Normalize((0.1307,), (0.3081,)),\n transforms.Lambda(shuffle_image)\n ]))\n\ntest_dataset = datasets.MNIST(root='./data',\n download=False,\n train=False,\n transform=transforms.Compose([\n transforms.ToTensor(),\n transforms.Normalize((0.1307,), (0.3081,)),\n transforms.Lambda(shuffle_image)\n ]))\n\ntrain_loader = DataLoader(train_dataset,\n batch_size=1,\n shuffle=True)\n\ntest_loader = DataLoader(test_dataset,\n batch_size=1,\n shuffle=True)\n\n\nclass MyModel(nn.Module):\n def __init__(self):\n super(MyModel, self).__init__()\n self.act = nn.ReLU()\n self.conv1 = nn.Conv2d(1, 4, 3, 1, 1)\n self.pool1 = nn.MaxPool2d(2)\n self.conv2 = nn.Conv2d(4, 8, 3, 1, 1)\n self.pool2 = nn.MaxPool2d(2)\n self.fc1 = nn.Linear(7*7*8, 10)\n\n def forward(self, x):\n x = self.act(self.conv1(x))\n x = self.pool1(x)\n x = self.act(self.conv2(x))\n x = self.pool2(x)\n x = x.view(x.size(0), -1)\n x = F.log_softmax(self.fc1(x), dim=1)\n return x\n\n\ndef train():\n acc = 0.0\n for batch_idx, (data, target) in enumerate(train_loader):\n optimizer.zero_grad()\n output = model(data)\n loss = criterion(output, target)\n loss.backward()\n optimizer.step()\n\n _, pred = torch.max(output, dim=1)\n accuracy = (pred == target).sum() / float(pred.size(0))\n acc += accuracy.data.float()\n\n if (batch_idx + 1) % 10 == 0:\n print('batch idx {}, loss {}'.format(\n batch_idx, loss.item()))\n\n acc /= len(train_loader)\n print('Train accuracy {}'.format(acc))\n\n\ndef test():\n acc = 0.0\n losses = 0.0\n for batch_idx, (data, target) in enumerate(test_loader):\n with torch.no_grad():\n output = model(data)\n loss = criterion(output, target)\n _, pred = torch.max(output, dim=1)\n\n accuracy = (pred == target).sum() / float(pred.size(0))\n acc += accuracy.data.float()\n losses += loss.item()\n acc /= len(test_loader)\n losses /= len(test_loader)\n print('Acc {}, loss {}'.format(\n acc, losses))\n\n\nmodel = MyModel()\ncriterion = nn.NLLLoss()\noptimizer = optim.Adam(model.parameters(), lr=1e-3)\n\ntrain()\ntest()\n\n# Visualize filters\nimport matplotlib.pyplot as plt\nfrom torchvision.utils import make_grid\nfilts1 = model.conv1.weight.data\ngrid = make_grid(filts1)\ngrid = grid.permute(1, 2, 0)\nplt.imshow(grid)\n" }, { "path": "model_sharding_data_parallel.py", "content": "\"\"\"\nModel sharding with DataParallel using 2 pairs of 2 GPUs.\n\n@author: ptrblck\n\"\"\"\n\nimport torch\nimport torch.nn as nn\n\n\nclass SubModule(nn.Module):\n def __init__(self, in_channels, out_channels):\n super(SubModule, self).__init__()\n self.conv1 = nn.Conv2d(in_channels, out_channels, 3, 1, 1)\n\n def forward(self, x):\n print('SubModule, device: {}, shape: {}\\n'.format(x.device, x.shape))\n x = self.conv1(x)\n return x\n\n\nclass MyModel(nn.Module):\n def __init__(self, split_gpus, parallel):\n super(MyModel, self).__init__()\n self.module1 = SubModule(3, 6)\n self.module2 = SubModule(6, 1)\n\n self.split_gpus = split_gpus\n self.parallel = parallel\n if self.split_gpus and self.parallel:\n self.module1 = nn.DataParallel(self.module1, device_ids=[0, 1]).to('cuda:0')\n self.module2 = nn.DataParallel(self.module2, device_ids=[2, 3]).to('cuda:2')\n\n def forward(self, x):\n print('Input: device {}, shape {}\\n'.format(x.device, x.shape))\n x = self.module1(x)\n print('After module1: device {}, shape {}\\n'.format(x.device, x.shape))\n x = self.module2(x)\n print('After module2: device {}, shape {}\\n'.format(x.device, x.shape))\n return x\n\n\nmodel = MyModel(split_gpus=True, parallel=True)\nx = torch.randn(16, 3, 24, 24).to('cuda:0')\noutput = model(x)\n" }, { "path": "momentum_update_nograd.py", "content": "\"\"\"\nScript to see how parameters are updated when an optimizer is used with\nmomentum/running estimates, even if gradients are zero.\n\nSet use_adam=True to see the effect. Otherwise plain SGD will be used.\n\nThe model consists of two \"decoder\" parts, dec1 and dec2.\nIn the first part of the script, you'll see that dec1 will be updated twice,\neven though this module is not used in the second forward pass.\nThis effect is observed, if one optimizer is used for all parameters.\n\nIn the second part of the script, two separate optimizers are used and\nwe cannot observe this effect anymore.\n\n@author: ptrblck\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\n\n\nuse_adam = True\n\n\nclass MyModel(nn.Module):\n def __init__(self):\n super(MyModel, self).__init__()\n self.enc = nn.Linear(64, 10)\n self.dec1 = nn.Linear(10, 64)\n self.dec2 = nn.Linear(10, 64)\n\n def forward(self, x, decoder_idx):\n x = F.relu(self.enc(x))\n if decoder_idx == 1:\n print('Using dec1')\n x = self.dec1(x)\n elif decoder_idx == 2:\n print('Using dec2')\n x = self.dec2(x)\n else:\n print('Unknown decoder_idx')\n\n return x\n\n\n# Create input and model\nx = torch.randn(1, 64)\ny = x.clone()\nmodel = MyModel()\ncriterion = nn.MSELoss()\n# Create optimizer using all model parameters\nif use_adam:\n optimizer = optim.Adam(model.parameters(), lr=1.)\nelse:\n optimizer = optim.SGD(model.parameters(), lr=1.)\n\n# Save init values\nold_state_dict = {}\nfor key in model.state_dict():\n old_state_dict[key] = model.state_dict()[key].clone()\n\n# Training procedure\noptimizer.zero_grad()\noutput = model(x, 1)\nloss = criterion(output, y)\nloss.backward()\n\n# Check for gradients in dec1, dec2\nprint('Dec1 grad: {}\\nDec2 grad: {}'.format(\n model.dec1.weight.grad, model.dec2.weight.grad))\n\noptimizer.step()\n\n# Save new params\nnew_state_dict = {}\nfor key in model.state_dict():\n new_state_dict[key] = model.state_dict()[key].clone()\n\n# Compare params\nfor key in old_state_dict:\n if not (old_state_dict[key] == new_state_dict[key]).all():\n print('Diff in {}'.format(key))\n\n# Update\nold_state_dict = {}\nfor key in model.state_dict():\n old_state_dict[key] = model.state_dict()[key].clone()\n\n# Pass through dec2\noptimizer.zero_grad()\noutput = model(x, 2)\nloss = criterion(output, y)\nloss.backward()\n\nprint('Dec1 grad: {}\\nDec2 grad: {}'.format(\n model.dec1.weight.grad, model.dec2.weight.grad))\n\noptimizer.step()\n\n# Save new params\nnew_state_dict = {}\nfor key in model.state_dict():\n new_state_dict[key] = model.state_dict()[key].clone()\n\n# Compare params\nfor key in old_state_dict:\n if not (old_state_dict[key] == new_state_dict[key]).all():\n print('Diff in {}'.format(key))\n\n## Create separate optimizers\nmodel = MyModel()\ndec1_params = list(model.enc.parameters()) + list(model.dec1.parameters())\noptimizer1 = optim.Adam(dec1_params, lr=1.)\ndec2_params = list(model.enc.parameters()) + list(model.dec2.parameters())\noptimizer2 = optim.Adam(dec2_params, lr=1.)\n\n# Save init values\nold_state_dict = {}\nfor key in model.state_dict():\n old_state_dict[key] = model.state_dict()[key].clone()\n\n# Training procedure\noptimizer1.zero_grad()\noutput = model(x, 1)\nloss = criterion(output, y)\nloss.backward()\n\n# Check for gradients in dec1, dec2\nprint('Dec1 grad: {}\\nDec2 grad: {}'.format(\n model.dec1.weight.grad, model.dec2.weight.grad))\n\noptimizer1.step()\n\n# Save new params\nnew_state_dict = {}\nfor key in model.state_dict():\n new_state_dict[key] = model.state_dict()[key].clone()\n\n# Compare params\nfor key in old_state_dict:\n if not (old_state_dict[key] == new_state_dict[key]).all():\n print('Diff in {}'.format(key))\n\n# Update\nold_state_dict = {}\nfor key in model.state_dict():\n old_state_dict[key] = model.state_dict()[key].clone()\n\n# Pass through dec2\noptimizer1.zero_grad()\noutput = model(x, 2)\nloss = criterion(output, y)\nloss.backward()\n\nprint('Dec1 grad: {}\\nDec2 grad: {}'.format(\n model.dec1.weight.grad, model.dec2.weight.grad))\n\noptimizer2.step()\n\n# Save new params\nnew_state_dict = {}\nfor key in model.state_dict():\n new_state_dict[key] = model.state_dict()[key].clone()\n\n# Compare params\nfor key in old_state_dict:\n if not (old_state_dict[key] == new_state_dict[key]).all():\n print('Diff in {}'.format(key))\n" }, { "path": "pytorch_redis.py", "content": "\"\"\"\nShows how to store and load data from redis using a PyTorch\nDataset and DataLoader (with multiple workers).\n\n@author: ptrblck\n\"\"\"\n\nimport redis\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.transforms as transforms\n\nimport numpy as np\n\n\n# Create random data and push to redis\nr = redis.Redis(host='localhost', port=6379, db=0)\n\nnb_images = 100\nfor idx in range(nb_images):\n # Use long for the fake images, as it's easier to store the target with it\n data = np.random.randint(0, 256, (3, 24, 24), dtype=np.long).tobytes()\n target = bytes(np.random.randint(0, 10, (1,)).astype(np.long))\n r.set(idx, data + target)\n\n\n# Create RedisDataset\nclass RedisDataset(Dataset):\n def __init__(self,\n redis_host='localhost',\n redis_port=6379,\n redis_db=0,\n length=0,\n transform=None):\n\n self.db = redis.Redis(host=redis_host, port=redis_port, db=redis_db)\n self.length = length\n self.transform = transform\n\n def __getitem__(self, index):\n data = self.db.get(index)\n data = np.frombuffer(data, dtype=np.long)\n x = data[:-1].reshape(3, 24, 24).astype(np.uint8)\n y = torch.tensor(data[-1]).long()\n if self.transform:\n x = self.transform(x)\n\n return x, y\n\n def __len__(self):\n return self.length\n\n\n# Load samples from redis using multiprocessing\ndataset = RedisDataset(length=100, transform=transforms.ToTensor())\nloader = DataLoader(\n dataset,\n batch_size=10,\n num_workers=2,\n shuffle=True\n)\n\nfor data, target in loader:\n print(data.shape)\n print(target.shape)\n" }, { "path": "shared_array.py", "content": "\"\"\"\nScript to demonstrate the usage of shared arrays using multiple workers.\n\nIn the first epoch the shared arrays in the dataset will be filled with\nrandom values. After setting set_use_cache(True), the shared values will be\nloaded from multiple processes.\n\n@author: ptrblck\n\"\"\"\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\n\nimport ctypes\nimport multiprocessing as mp\n\nimport numpy as np\n\n\nclass MyDataset(Dataset):\n def __init__(self):\n shared_array_base = mp.Array(ctypes.c_float, nb_samples*c*h*w)\n shared_array = np.ctypeslib.as_array(shared_array_base.get_obj())\n shared_array = shared_array.reshape(nb_samples, c, h, w)\n self.shared_array = torch.from_numpy(shared_array)\n self.use_cache = False\n\n def set_use_cache(self, use_cache):\n self.use_cache = use_cache\n\n def __getitem__(self, index):\n if not self.use_cache:\n print('Filling cache for index {}'.format(index))\n # Add your loading logic here\n self.shared_array[index] = torch.randn(c, h, w)\n x = self.shared_array[index]\n return x\n\n def __len__(self):\n return nb_samples\n\n\nnb_samples, c, h, w = 10, 3, 24, 24\n\ndataset = MyDataset()\nloader = DataLoader(\n dataset,\n num_workers=2,\n shuffle=False\n)\n\nfor epoch in range(2):\n for idx, data in enumerate(loader):\n print('Epoch {}, idx {}, data.shape {}'.format(epoch, idx, data.shape))\n\n if epoch == 0:\n loader.dataset.set_use_cache(True)\n" }, { "path": "shared_dict.py", "content": "\"\"\"\nScript to demonstrate the usage of shared dicts using multiple workers.\n\nIn the first epoch the shared dict in the dataset will be filled with\nrandom values. The next epochs will just use the dict without \"loading\" the\ndata again.\n\n@author: ptrblck\n\"\"\"\n\nfrom multiprocessing import Manager\n\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\n\n\nclass MyDataset(Dataset):\n def __init__(self, shared_dict, length):\n self.shared_dict = shared_dict\n self.length = length\n\n def __getitem__(self, index):\n if index not in self.shared_dict:\n print('Adding {} to shared_dict'.format(index))\n self.shared_dict[index] = torch.tensor(index)\n return self.shared_dict[index]\n\n def __len__(self):\n return self.length\n\n\n# Init\nmanager = Manager()\nshared_dict = manager.dict()\ndataset = MyDataset(shared_dict, length=100)\n\nloader = DataLoader(\n dataset,\n batch_size=10,\n num_workers=6,\n shuffle=True,\n pin_memory=True\n)\n\n# First loop will add data to the shared_dict\nfor x in loader:\n print(x)\n\n# The second loop will just get the data\nfor x in loader:\n print(x)\n" }, { "path": "unet_demo.py", "content": "\"\"\"\nSimple UNet demo\n\n@author: ptrblck\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\n\n\nclass BaseConv(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, padding,\n stride):\n super(BaseConv, self).__init__()\n\n self.act = nn.ReLU()\n\n self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size, padding,\n stride)\n\n self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size,\n padding, stride)\n\n def forward(self, x):\n x = self.act(self.conv1(x))\n x = self.act(self.conv2(x))\n return x\n\n\nclass DownConv(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, padding,\n stride):\n super(DownConv, self).__init__()\n\n self.pool1 = nn.MaxPool2d(kernel_size=2)\n self.conv_block = BaseConv(in_channels, out_channels, kernel_size,\n padding, stride)\n\n def forward(self, x):\n x = self.pool1(x)\n x = self.conv_block(x)\n return x\n\n\nclass UpConv(nn.Module):\n def __init__(self, in_channels, in_channels_skip, out_channels,\n kernel_size, padding, stride):\n super(UpConv, self).__init__()\n\n self.conv_trans1 = nn.ConvTranspose2d(\n in_channels, in_channels, kernel_size=2, padding=0, stride=2)\n self.conv_block = BaseConv(\n in_channels=in_channels + in_channels_skip,\n out_channels=out_channels,\n kernel_size=kernel_size,\n padding=padding,\n stride=stride)\n\n def forward(self, x, x_skip):\n x = self.conv_trans1(x)\n x = torch.cat((x, x_skip), dim=1)\n x = self.conv_block(x)\n return x\n\n\nclass UNet(nn.Module):\n def __init__(self, in_channels, out_channels, n_class, kernel_size,\n padding, stride):\n super(UNet, self).__init__()\n\n self.init_conv = BaseConv(in_channels, out_channels, kernel_size,\n padding, stride)\n\n self.down1 = DownConv(out_channels, 2 * out_channels, kernel_size,\n padding, stride)\n\n self.down2 = DownConv(2 * out_channels, 4 * out_channels, kernel_size,\n padding, stride)\n\n self.down3 = DownConv(4 * out_channels, 8 * out_channels, kernel_size,\n padding, stride)\n\n self.up3 = UpConv(8 * out_channels, 4 * out_channels, 4 * out_channels,\n kernel_size, padding, stride)\n\n self.up2 = UpConv(4 * out_channels, 2 * out_channels, 2 * out_channels,\n kernel_size, padding, stride)\n\n self.up1 = UpConv(2 * out_channels, out_channels, out_channels,\n kernel_size, padding, stride)\n\n self.out = nn.Conv2d(out_channels, n_class, kernel_size, padding, stride)\n\n def forward(self, x):\n # Encoder\n x = self.init_conv(x)\n x1 = self.down1(x)\n x2 = self.down2(x1)\n x3 = self.down3(x2)\n # Decoder\n x_up = self.up3(x3, x2)\n x_up = self.up2(x_up, x1)\n x_up = self.up1(x_up, x)\n x_out = F.log_softmax(self.out(x_up), 1)\n return x_out\n\n\n# Create 10-class segmentation dummy image and target\nnb_classes = 10\nx = torch.randn(1, 3, 96, 96)\ny = torch.randint(0, nb_classes, (1, 96, 96))\n\nmodel = UNet(in_channels=3,\n out_channels=64,\n n_class=10,\n kernel_size=3,\n padding=1,\n stride=1)\n\nif torch.cuda.is_available():\n model = model.to('cuda')\n x = x.to('cuda')\n y = y.to('cuda')\n\ncriterion = nn.NLLLoss()\noptimizer = optim.Adam(model.parameters(), lr=1e-3)\n\n# Training loop\nfor epoch in range(1):\n optimizer.zero_grad()\n\n output = model(x)\n loss = criterion(output, y)\n loss.backward()\n optimizer.step()\n\n print('Epoch {}, Loss {}'.format(epoch, loss.item()))\n" }, { "path": "weighted_sampling.py", "content": "\"\"\"\nUsage of WeightedRandomSampler using an imbalanced dataset with\nclass imbalance 99 to 1.\n\n@author: ptrblck\n\"\"\"\n\nimport torch\nfrom torch.utils.data.sampler import WeightedRandomSampler\nfrom torch.utils.data.dataloader import DataLoader\n\n\n# Create dummy data with class imbalance 99 to 1\nnumDataPoints = 1000\ndata_dim = 5\nbs = 100\ndata = torch.randn(numDataPoints, data_dim)\ntarget = torch.cat((torch.zeros(int(numDataPoints * 0.99), dtype=torch.long),\n torch.ones(int(numDataPoints * 0.01), dtype=torch.long)))\n\nprint('target train 0/1: {}/{}'.format(\n (target == 0).sum(), (target == 1).sum()))\n\n# Compute samples weight (each sample should get its own weight)\nclass_sample_count = torch.tensor(\n [(target == t).sum() for t in torch.unique(target, sorted=True)])\nweight = 1. / class_sample_count.float()\nsamples_weight = torch.tensor([weight[t] for t in target])\n\n# Create sampler, dataset, loader\nsampler = WeightedRandomSampler(samples_weight, len(samples_weight))\ntrain_dataset = torch.utils.data.TensorDataset(data, target)\ntrain_loader = DataLoader(\n train_dataset, batch_size=bs, num_workers=1, sampler=sampler)\n\n# Iterate DataLoader and check class balance for each batch\nfor i, (x, y) in enumerate(train_loader):\n print(\"batch index {}, 0/1: {}/{}\".format(\n i, (y == 0).sum(), (y == 1).sum()))\n" } ]