Homec4science
Diffusion LPBF Transfer Learning (master) Browse

R11778/resnet.py
0a854839b33emaster

Last Change
Raw File

resnet.py
View Options

# -*- coding: utf-8 -*-
"""
Created on Sat Feb 8 22:10:18 2020
@author: srpv
"""
import torch
import numpy as np
import torchvision.transforms as transforms
import torchvision
import torch.nn.functional as F
class baseBlock(torch.nn.Module):
expansion = 1
def __init__(self,input_planes,planes,stride=1,dim_change=None):
super(baseBlock,self).__init__()
#declare convolutional layers with batch norms
self.conv1 = torch.nn.Conv2d(input_planes,planes,stride=stride,kernel_size=3,padding=1)
self.bn1 = torch.nn.BatchNorm2d(planes)
self.conv2 = torch.nn.Conv2d(planes,planes,stride=1,kernel_size=3,padding=1)
self.bn2 = torch.nn.BatchNorm2d(planes)
self.dim_change = dim_change
def forward(self,x):
#Save the residue
res = x
output = F.relu(self.bn1(self.conv1(x)))
output = self.bn2(self.conv2(output))
if self.dim_change is not None:
res = self.dim_change(res)
output += res
output = F.relu(output)
return output
class bottleNeck(torch.nn.Module):
expansion = 4
def __init__(self,input_planes,planes,stride=1,dim_change=None):
super(bottleNeck,self).__init__()
self.conv1 = torch.nn.Conv2d(input_planes,planes,kernel_size=1,stride=1)
self.bn1 = torch.nn.BatchNorm2d(planes)
self.conv2 = torch.nn.Conv2d(planes,planes,kernel_size=3,stride=stride,padding=1)
self.bn2 = torch.nn.BatchNorm2d(planes)
self.conv3 = torch.nn.Conv2d(planes,planes*self.expansion,kernel_size=1)
self.bn3 = torch.nn.BatchNorm2d(planes*self.expansion)
self.dim_change = dim_change
def forward(self,x):
res = x
output = F.relu(self.bn1(self.conv1(x)))
output = F.relu(self.bn2(self.conv2(output)))
output = self.bn3(self.conv3(output))
if self.dim_change is not None:
res = self.dim_change(res)
output += res
output = F.relu(output)
return output
class ResNet(torch.nn.Module):
def __init__(self,block,num_layers,classes=4):
super(ResNet,self).__init__()
#according to research paper:
self.input_planes = 64
self.conv1 = torch.nn.Conv2d(3,64,kernel_size=3,stride=1,padding=1)
self.bn1 = torch.nn.BatchNorm2d(64)
self.layer1 = self._layer(block,64,num_layers[0],stride=1)
self.layer2 = self._layer(block,128,num_layers[1],stride=2)
self.layer3 = self._layer(block,256,num_layers[2],stride=2)
self.layer4 = self._layer(block,512,num_layers[3],stride=2)
self.averagePool = torch.nn.AvgPool2d(kernel_size=4,stride=1)
self.fc = torch.nn.Linear(512*block.expansion,classes)
def _layer(self,block,planes,num_layers,stride=1):
dim_change = None
if stride!=1 or planes != self.input_planes*block.expansion:
dim_change = torch.nn.Sequential(torch.nn.Conv2d(self.input_planes,planes*block.expansion,kernel_size=1,stride=stride),
torch.nn.BatchNorm2d(planes*block.expansion))
netLayers =[]
netLayers.append(block(self.input_planes,planes,stride=stride,dim_change=dim_change))
self.input_planes = planes * block.expansion
for i in range(1,num_layers):
netLayers.append(block(self.input_planes,planes))
self.input_planes = planes * block.expansion
return torch.nn.Sequential(*netLayers)
def forward(self,x):
x = F.relu(self.bn1(self.conv1(x)))
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = F.avg_pool2d(x,4)
x = x.view(x.size(0),-1)
x = self.fc(x)
return x
c4science · Help