ISRLU_model_pytorch.py

# -*- coding: utf-8 -*-
"""Inverse Square Root Linear Unit.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/19PF-Q85uhsT9eYtsr0MkRtAZfXaCDL6w
"""

import numpy as np
import pandas as pd
from collections import OrderedDict
import math

import torch
from torch.autograd import Variable
import torch.nn as nn 
from torch.autograd import Function
from torch.nn.parameter import Parameter
from torch import optim
import torch.nn.functional as F
from torchvision import datasets,transforms

def train_model(model,train_loader):
  criterion = nn.NLLLoss()
  learning_rate = 0.003
  epochs = 5
  optimizer = optim.Adam(model.parameters() , lr = learning_rate)

  print('Training the model. Make sure that loss decreases after each epoch.\n')

  for e in range(epochs):
    running_loss = 0
    for images,labels in trainloader:
      images = images.view(images.shape[0], -1)
      log_ps = model(images)
      loss = criterion(log_ps , labels)

      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

      running_loss += loss.item()
    else:
       print(f"Training loss: {running_loss}")


def ISRLU(x):  

  if(x >= 0):
    return x

  else:
    alpha=3    
    den = math.sqrt(1 + alpha * (math.pow(x,2))
    return x / den



class ISRLU(nn.Module):
  def __init__(self):
    super().__init__()

  def forward(self,input):
    return ISRLU(x)


class ClassifierISRLU(nn.Module):
  def __init__(self):
        super().__init__()

        # initialize layers
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 64)
        self.fc4 = nn.Linear(64, 10)

  def forward(self, x):
      # make sure the input tensor is flattened
      x = x.view(x.shape[0], -1)

      # apply silu function
      x = ISRLU(self.fc1(x))

      # apply silu function
      x = ISRLU(self.fc2(x))

      # apply silu function
      x = ISRLU(self.fc3(x))

      x = F.log_softmax(self.fc4(x), dim=1)

      return x