import torch
from torch.autograd import Variable
import torch.nn as nn
import pickle

from sys import path
path.append('/residual model path')
import residual_model
from residual_model import Residual_Model

model = Residual_Model()
model.load_state_dict(torch.load('./model.pkl'))



class myNet(nn.Module):
  def __init__(self,pretrained_model,layers):
    super(myNet,self).__init__()
    self.net1 = nn.Sequential(*list(pretrained_model.children())[:layers[0]])
    self.net2 = nn.Sequential(*list(pretrained_model.children())[:layers[1]])
    self.net3 = nn.Sequential(*list(pretrained_model.children())[:layers[2]])

  def forward(self,x):
    out1 = self.net1(x)
    out2 = self.net(out1)
    out3 = self.net(out2)
    return out1,out2,out3

def get_features(pretrained_model, x, layers = [3, 4, 9]): ## get_features 其實很簡單
'''
1.首先import model 
2.將weights load 進model
3.熟悉model的每一層的位置，提前知道要輸出feature map的網絡層是處于網絡的那一層
4.直接將test_x輸入網絡，*list(model.chidren())是用來提取網絡的每一層的結構的。net1 = nn.Sequential(*list(pretrained_model.children())[:layers[0]]) ,就是第三層前的所有層。

'''
  net1 = nn.Sequential(*list(pretrained_model.children())[:layers[0]]) 
#  print net1 
  out1 = net1(x) 

  net2 = nn.Sequential(*list(pretrained_model.children())[layers[0]:layers[1]]) 
#  print net2 
  out2 = net2(out1) 

  #net3 = nn.Sequential(*list(pretrained_model.children())[layers[1]:layers[2]]) 
  #out3 = net3(out2) 

  return out1, out2
with open('test.pickle','rb') as f:
  data = pickle.load(f)
x = data['test_mains'][0]
x = Variable(torch.from_numpy(x)).view(1,1,128,1) ## test_x必須為Varibable
#x = Variable(torch.randn(1,1,128,1))
if torch.cuda.is_available():
  x = x.cuda() # 如果模型的訓練是用cuda加速的話，輸入的變量也必須是cuda加速的，兩個必須是對應的，網絡的參數weight都是用cuda加速的，不然會報錯
  model = model.cuda()
output1,output2 = get_features(model,x)## model是訓練好的model,前面已經import 進來了Residual model
print('output1.shape:',output1.shape)
print('output2.shape:',output2.shape)
#print('output3.shape:',output3.shape)
output_1 = torch.squeeze(output2,dim = 0)
output_1_arr = output_1.data.cpu().numpy() # 得到的cuda加速的輸出不能直接轉變成numpy格式的，當時根據報錯的信息首先將變量轉換為cpu的，然后轉換為numpy的格式
output_1_arr = output_1_arr.reshape([output_1_arr.shape[0],output_1_arr.shape[1]])