Spaces:
Sleeping
Sleeping
File size: 1,480 Bytes
2794d4e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
import torch
import torch.nn as nn
from torch import Tensor as tensor
from torch.nn import functional as F
import torchvision.models as models
class SCA_Block(nn.Module):
def __init__(self, in_channel, downsample_channel):
super().__init__()
self.conv_A = nn.Conv2d(in_channel, downsample_channel, (1,1))
self.conv_B = nn.Conv2d(in_channel, downsample_channel, (1,1))
self.conv_E = nn.Conv2d(in_channel, downsample_channel, (1,1))
self.linear = nn.Linear(downsample_channel,in_channel)
def forward(self, feature_in):
b_size,c,w,h = feature_in.shape
A = self.conv_A(feature_in)
B = self.conv_B(feature_in)
E = self.conv_E(feature_in)
c1 = A.shape[1]
Z = F.softmax(torch.dot(torch.reshape(A,(b_size,c1,-1)),
torch.reshape(B,(b_size,-1,c1))), axis = 1 )
D = torch.reshape( Z * torch.reshape(E,(b_size,c1,-1)) , (b_size,c1,w,h))
out = feature_in * F.sigmoid(F.adaptive_avg_pool2d(D))
return out
class AQC_NET(nn.Module):
def __init__(self, pretrain = True, num_label = 5):
super().__init__()
self.resnet18 = models.resnet18(pretrained = pretrain)
self.resnet18.layer3[0].add_module('sca_1', SCA_Block(256,16))
self.resnet18.layer3[1].add_module('sca_2', SCA_Block(256,16))
self.resnet18.fc = nn.Linear(512,num_label)
def forward(self,x):
return F.softmax(self.resnet18(x))
|