When I create neural software systems, I most often use the PyTorch library. The Keras library is very good for basic neural systems but for advanced architectures I like the flexibility of PyTorch. Using raw TensorFlow without Keras is an option, but I am more comfortable using the PyTorch APIs.


An example of a custom NoisyLinear() layer. Notice the two outputs are slightly different.

I hadn't looked at the problem of creating a custom PyTorch Layer in several months, so I figured I'd code up a demo. The most fundamental layer is Linear(). For a 4-7-3 neural network (four input nodes, one hidden layer with seven nodes, three output nodes), a definition could look like:

  import torch as T    class Net(T.nn.Module):    def __init__(self):      super(Net, self).__init__()      self.hid1 = T.nn.Linear(4, 7)  # 4-7-3      self.oupt = T.nn.Linear(7, 3)  # default init      def forward(self, x):      z = T.tanh(self.hid1(x))      z = self.oupt(z)      return z

For my demo, I decided to create a custom NoisyLinear() layer that works just like a standard Linear() layer but injects randomness. This isn't particularly useful by itself but I'm just experimenting. So I wanted a 4-7-3 network to work like this:

  class Net(T.nn.Module):    def __init__(self):      super(Net, self).__init__()      self.hid1 = NoisyLinear(4, 7)  # 4-7-3      self.oupt = NoisyLinear(7, 3)      def forward(self, x):      z = T.tanh(self.hid1(x))      z = self.oupt(z)       return z

In other words, everything is the same except I use the program defined NoisyLinear() instead of the built-in torch.nn.Linear() layer. The custom layer definition I came up with is:

  class NoisyLinear(T.nn.Module):    def __init__(self, n_in, n_out):      super().__init__()      self.n_in, self.n_out = n_in, n_out        self.weights = T.nn.Parameter(T.zeros((n_out, n_in),        dtype=T.float32))      self.bias = T.nn.Parameter(T.tensor(n_out,        dtype=T.float32))      self.lo = 0.90; self.hi = 0.98  # noise        lim = 0.01  # initialize weights and bias      T.nn.init.uniform_(self.weights, -lim, +lim)      T.nn.init.uniform_(self.bias, -lim, +lim)      def forward(self, x):      wx= T.mm(x, self.weights.t())      rnd = (self.hi - self.lo) * T.rand(1) + self.lo      return rnd * T.add(wx, self.bias)  # wts * x + bias

The Parameter() class makes the weights and the bias trainable. I used basic uniform initialization with hard-coded range [-0.01, +0.01]. The forward() method computes weights * inputs + bias as usual, but then multiples the results by random noise in the range [0.90, 0.98]. Each time the forward() method of a NoisyLayer() layer instance is called, the result will be slightly different.

Writing a custom layer for PyTorch is rarely needed, but compared to alternative libraries, customizing PyTorch is relatively easier -- with an emphasis on "relatively".


Three well-known custom cars. Left: Dodge Deodora (1965). Center: Norman Timbs Special (1947). Right: Chrysler Thunderbolt (1941).

Complete demo code below. Long.

  # iris_noisy_layer.py  # creating a custom "NoisyLinear" layer  # PyTorch 1.9.0-CPU Anaconda3-2020.02  Python 3.7.6  # Windows 10     import numpy as np  import torch as T    device = T.device("cpu")  # to Tensor or Module    # -----------------------------------------------------------    class NoisyLinear(T.nn.Module):    def __init__(self, n_in, n_out):      super().__init__()      self.n_in, self.n_out = n_in, n_out        self.weights = T.nn.Parameter(T.zeros((n_out, n_in),        dtype=T.float32))      self.bias = T.nn.Parameter(T.tensor(n_out,        dtype=T.float32))      self.lo = 0.90; self.hi = 0.98  # noise        lim = 0.01  # initialize weights and bias      T.nn.init.uniform_(self.weights, -lim, +lim)      T.nn.init.uniform_(self.bias, -lim, +lim)      def forward(self, x):      wx= T.mm(x, self.weights.t())      rnd = (self.hi - self.lo) * T.rand(1) + self.lo      return rnd * T.add(wx, self.bias)  # wts * x + bias    # -----------------------------------------------------------    class IrisDataset(T.utils.data.Dataset):    def __init__(self, src_file, num_rows=None):      # 5.0, 3.5, 1.3, 0.3, 0      tmp_x = np.loadtxt(src_file, max_rows=num_rows,        usecols=range(0,4), delimiter=",", skiprows=0,        dtype=np.float32)      tmp_y = np.loadtxt(src_file, max_rows=num_rows,        usecols=4, delimiter=",", skiprows=0,        dtype=np.int64)        self.x_data = T.tensor(tmp_x, dtype=T.float32)      self.y_data = T.tensor(tmp_y, dtype=T.int64)      def __len__(self):      return len(self.x_data)      def __getitem__(self, idx):      if T.is_tensor(idx):        idx = idx.tolist()      preds = self.x_data[idx]      spcs = self.y_data[idx]       sample = { 'predictors' : preds, 'species' : spcs }      return sample    # -----------------------------------------------------------    class Net(T.nn.Module):    def __init__(self):      super(Net, self).__init__()      self.hid1 = NoisyLinear(4, 7)  # 4-7-3      self.oupt = NoisyLinear(7, 3)      def forward(self, x):      z = T.tanh(self.hid1(x))      z = self.oupt(z)  # no softmax: CrossEntropyLoss()       return z    # -----------------------------------------------------------    def accuracy(model, dataset):    # assumes model.eval()    dataldr = T.utils.data.DataLoader(dataset, batch_size=1,      shuffle=False)    n_correct = 0; n_wrong = 0    for (_, batch) in enumerate(dataldr):      X = batch['predictors']       # Y = T.flatten(batch['species'])      Y = batch['species']  # already flattened by Dataset      with T.no_grad():        oupt = model(X)  # logits form        big_idx = T.argmax(oupt)      # if big_idx.item() == Y.item():      if big_idx == Y:        n_correct += 1      else:        n_wrong += 1      acc = (n_correct * 1.0) / (n_correct + n_wrong)    return acc    # -----------------------------------------------------------    def main():    # 0. get started    print("\nBegin Iris custom NoisyLinear layer demo \n")    T.manual_seed(1)    np.random.seed(1)        # 1. create Dataset and DataLoader objects    print("Creating Iris train DataLoader ")      train_file = ".\\Data\\iris_train.txt"    train_ds = IrisDataset(train_file, num_rows=120)      bat_size = 4    train_ldr = T.utils.data.DataLoader(train_ds,      batch_size=bat_size, shuffle=True)      # 2. create network    net = Net().to(device)      # 3. train model    max_epochs = 20    ep_log_interval = 4    lrn_rate = 0.05      loss_func = T.nn.CrossEntropyLoss()  # applies softmax()    optimizer = T.optim.SGD(net.parameters(), lr=lrn_rate)      print("\nbat_size = %3d " % bat_size)    print("loss = " + str(loss_func))    print("optimizer = SGD")    print("max_epochs = %3d " % max_epochs)    print("lrn_rate = %0.3f " % lrn_rate)      print("\nStarting training")    net.train()    for epoch in range(0, max_epochs):      epoch_loss = 0  # for one full epoch      num_lines_read = 0        for (batch_idx, batch) in enumerate(train_ldr):        X = batch['predictors']  # [10,4]        Y = batch['species']  # OK; alreay flattened          optimizer.zero_grad()        oupt = net(X)        loss_val = loss_func(oupt, Y)  # a tensor        epoch_loss += loss_val.item()  # accumulate        loss_val.backward()            # gradients        optimizer.step()               # update wts        if epoch % ep_log_interval == 0:        print("epoch = %4d   loss = %0.4f" % (epoch, epoch_loss))    print("Done ")      # 4. evaluate model accuracy    print("\nComputing model accuracy")    net.eval()    acc = accuracy(net, train_ds)  # item-by-item    print("Accuracy on train data = %0.4f" % acc)      # 5. make a prediction    print("\nPredicting species for [6.1, 3.1, 5.1, 1.1]: ")    x = np.array([[6.1, 3.1, 5.1, 1.1]], dtype=np.float32)    x = T.tensor(x, dtype=T.float32).to(device)       with T.no_grad():      logits = net(x).to(device)  # values do not sum to 1.0    probs = T.softmax(logits, dim=1).to(device)    T.set_printoptions(precision=4)    print(probs)      print("\nPredicting again for [6.1, 3.1, 5.1, 1.1]: ")    x = np.array([[6.1, 3.1, 5.1, 1.1]], dtype=np.float32)    x = T.tensor(x, dtype=T.float32).to(device)       with T.no_grad():      logits = net(x).to(device)  # values do not sum to 1.0    probs = T.softmax(logits, dim=1).to(device)    T.set_printoptions(precision=4)    print(probs)      print("\nEnd custom NoisyLinear layer demo")    if __name__ == "__main__":    main()