first

network_multi.py

@@ -0,0 +1,37 @@
+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])