text: """
  Architecture definition for Fully Convolutional Neural Networks (FCN8s)
  Not initialised with pretrained VGG weights
  Instead weights initialised by Xavier initialisation for convolutional layers
  and bilinear upsampling for convolutional transpose layers
  """
  import torch
  import torch.nn as nn
  import torch.optim as optim
  import numpy as np
  import math
  import os
  import time
  
  #os.environ["CUDA_VISIBLE_DEVICES"]= "4"
  
  use_gpu = torch.cuda.is_available()
  num_gpu = list(range(torch.cuda.device_count()))
  
  	
  def get_upsampling_weight(in_channels,out_channels,kernel_size):
  	""" Make a 2D bilinear kernel suitable for upsampling"""
  	factor = (kernel_size+1)//2
  	if kernel_size%2 == 1:
  		center = factor-1
  	else:
  		center = factor-1
  	og = np.ogrid[:kernel_size, :kernel_size]
  	filt = (1 - abs(og[0] - center)/factor) * (1 - abs(og[1] - center)/factor)
  	weight = np.zeros((in_channels,out_channels,kernel_size,kernel_size), dtype = np.float64)
  	weight[range(in_channels), range(out_channels),:,: ]
  	return torch.from_numpy(weight).float()
  	
  class FCN8s(nn.Module):
  	def __init__(self,n_class): 
  		
  		super(FCN8s,self).__init__()                                #VGG architecture definition
  		#conv1
  		self.conv1_1 = nn.Conv2d(1,64,3,padding=100)
  		self.relu1_1 = nn.ReLU(inplace=True)
  		self.conv1_2 = nn.Conv2d(64,64,3,padding = 1)
  		self.relu1_2 = nn.ReLU(inplace=True)
  		self.pool1 = nn.MaxPool2d(2, stride=2, ceil_mode=True) #1/2 dimension reduction
  		
  		#conv2
  		self.conv2_1 = nn.Conv2d(64,128,3,padding = 1)
  		self.relu2_1 = nn.ReLU(inplace=True)
  		self.conv2_2 = nn.Conv2d(128,128,3,padding=1)
  		self.relu2_2 = nn.ReLU(inplace=True)
  		self.pool2 = nn.MaxPool2d(2, stride=2, ceil_mode=True) # 1/4 dimension reduction 
  		
  		
  		#conv3 
  		self.conv3_1 = nn.Conv2d(128,256,3,padding=1)
  		self.relu3_1 = nn.ReLU(inplace = True)
  		self.conv3_2 = nn.Conv2d(256,256,3,padding=1)
  		self.relu3_2 = nn.ReLU(inplace = True)
  		self.conv3_3 = nn.Conv2d(256,256,3,padding=1)
  		self.relu3_3 = nn.ReLU(inplace = True)
  		self.pool3 = nn.MaxPool2d(2, stride=2, ceil_mode=True)  # 1/8 dimension reduction
  		
  		#conv4
  		self.conv4_1 = nn.Conv2d(256, 512, 3, padding=1)
  		self.relu4_1 = nn.ReLU(inplace=True)
  		self.conv4_2 = nn.Conv2d(512, 512, 3, padding=1)
  		self.relu4_2 = nn.ReLU(inplace=True)
  		self.conv4_3 = nn.Conv2d(512, 512, 3, padding=1)
  		self.relu4_3 = nn.ReLU(inplace=True)
  		self.pool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True) # 1/16 dimension reduction
  		
  		#conv5
  		self.conv5_1 = nn.Conv2d(512, 512, 3, padding=1)
  		self.relu5_1 = nn.ReLU(inplace=True)
  		self.conv5_2 = nn.Conv2d(512, 512, 3, padding=1)
  		self.relu5_2 = nn.ReLU(inplace=True)
  		self.conv5_3 = nn.Conv2d(512, 512, 3, padding=1)
  		self.relu5_3 = nn.ReLU(inplace=True)
  		self.pool5 = nn.MaxPool2d(2, stride=2, ceil_mode=True) # 1/32 dimension reduction
  		
  		#fc6
  		self.fc6 = nn.Conv2d(512,4096,7)
  		self.relu6 = nn.ReLU(inplace=True)
  		self.drop6 = nn.Dropout2d()
  		
  		#fc7
  		self.fc7 = nn.Conv2d(4096,4096,1)
  		self.relu7 = nn.ReLU(inplace = True)
  		self.drop7 = nn.Dropout2d()
  		
  		self.score_fr = nn.Conv2d(4096,n_class,1)             #Skip Layer defintions
  		self.score_pool3 = nn.Conv2d(256,n_class,1)
  		self.score_pool4 = nn.Conv2d(512,n_class,1)
  		
  		self.upscore2 = nn.ConvTranspose2d(n_class,n_class,4,stride=2,bias=False)      #Upsampling layer defintions
  		self.upscore8 = nn.ConvTranspose2d(n_class,n_class,16,stride=8,bias=False)
  		self.upscore_pool4 = nn.ConvTranspose2d(n_class,n_class,4,stride=2,bias=False)
  		
  		self._initialize_weights()
  		
  	def _initialize_weights(self):
  		#Xavier initialisation
  		for m in self.modules():
  			if isinstance(m,nn.Conv2d):
  				n = m.kernel_size[0] * m.kernel_size[1] *m.out_channels
  				m.weight.data.normal_(0,math.sqrt(2./n))
  				#m.weight.data.zero_()
  				if m.bias is not None:
  						m.bias.data.zero_()           
  			if isinstance(m,nn.ConvTranspose2d):
  				assert m.kernel_size[0] == m.kernel_size[1]
  				initial_weight = get_upsampling_weight(m.in_channels,m.out_channels,m.kernel_size[0])
  				m.weight.data.copy_(initial_weight)
  					
  	def forward(self,x):
  		h = x
  		h = self.relu1_1(self.conv1_1(h))
  		h = self.relu1_2(self.conv1_2(h))
  		h = self.pool1(h)
  			
  		h = self.relu2_1(self.conv2_1(h))
  		h = self.relu2_2(self.conv2_2(h))
  		h = self.pool2(h)
  			
  		h = self.relu3_1(self.conv3_1(h))
  		h = self.relu3_2(self.conv3_2(h))
  		h = self.relu3_3(self.conv3_3(h))
  		h = self.pool3(h)
  		pool3 = h # 1/8
  			
  		h = self.relu4_1(self.conv4_1(h))
  		h = self.relu4_2(self.conv4_2(h))
  		h = self.relu4_3(self.conv4_3(h))
  		h = self.pool4(h)
  		pool4 = h # 1/16
  			
  		h = self.relu5_1(self.conv5_1(h))
  		h = self.relu5_2(self.conv5_2(h))
  		h = self.relu5_3(self.conv5_3(h))
  		h = self.pool5(h)
  			
  		h = self.relu6(self.fc6(h))
  		h = self.drop6(h)
  			
  		h = self.relu7(self.fc7(h))
  		h = self.drop7(h)
  			
  		h = self.score_fr(h)
  		h = self.upscore2(h)
  		upscore2 = h  # 1/16
  			
  		h = self.score_pool4(pool4)
  		h = h[:, :, 5:5 + upscore2.size()[2], 5:5 + upscore2.size()[3]]
  		score_pool4c = h
  			
  		h = upscore2 + score_pool4c # 1/16
  		h = self.upscore_pool4(h)
  		upscore_pool4 = h # 1/8
  			
  		h = self.score_pool3(pool3)
  		h = h[:,:,9:9+upscore_pool4.size()[2],9:9+upscore_pool4.size()[3]]
  		score_pool3c = h
  			
  		h = upscore_pool4 + score_pool3c
  			
  		h = self.upscore8(h)
  		h = h[:, :, 31:31 + x.size()[2], 31:31 + x.size()[3]].contiguous()
  		
  			
  		return h  	
  			
  """
  #Test code
  if __name__ == "__main__":
  	batch_size, n_class, h,w = 5,11,224,224
  	
  	#test the output size
  	
  	fcn_model = FCN8s(n_class)
  	if use_gpu:
  		ts = time.time()
  		fcn_model = fcn_model.cuda()
  		print ("Finsished loading CUDA, time elapsed {}".format(time.time()-ts))
  	
  	input = torch.autograd.Variable(torch.randn(batch_size,3,h,w).cuda())
  	print("hello")
  	output = fcn_model(input)
  	print(output.size())
  	
  	#To check whether training properly
  	y = torch.autograd.Variable(torch.randn(batch_size,n_class,h,w).cuda())
  	criterion = nn.BCEWithLogitsLoss()
  	optimiser = optim.SGD(fcn_model.parameters(),lr=1e-3,momentum=0.9)
  	for iter in range(10):
  		optimiser.zero_grad()
  		output = fcn_model(input)
  		loss = criterion(output,y)
  		loss.backward()
  		print("iter {}, loss {}".format(iter, loss.data[0]))
  		optimiser.step()
  """