38 lines
1010 B
Python
Executable File
38 lines
1010 B
Python
Executable File
import torch.nn as nn
|
|
import torch
|
|
#from pytorch_pretrained_vit import ViT,VoT,VVT
|
|
from vgg import vgg16 as backbone
|
|
from common import global_para as gp
|
|
|
|
|
|
|
|
class Multi_transform_image(nn.Module):
|
|
def __init__(self, load_pretrain=True):
|
|
super(Multi_transform_image, self).__init__()
|
|
#frontend_feat
|
|
#self.frontend = backbone(pretrained=load_pretrain)
|
|
self.pixel_model = ViT('B_16_imagenet1k', pretrained=True,num_classes=gp.FineClassNum,image_size=448)
|
|
def forward(self, x):
|
|
x = x.float()
|
|
output_pixel = self.pixel_model(x)
|
|
return output_pixel
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
|
|
print(device)
|
|
net = Multi_transform().to(device)
|
|
input1 = torch.randn(1, 3, 448,448).to(device)
|
|
input2 = torch.randn(1, 3, 448, 448).to(device)
|
|
output = net(input1,input2)
|
|
print(output.shape)
|
|
# torch.Size([1, 1, 96, 128]) torch.Size([1, 5]) torch.Size([1, 13])
|