maml.py

# -*- coding: utf-8 -*-#

#-------------------------------------------------------------------------------
# Name:         maml
# Description:  simple maml with 4 conv layers
# Author:       Boliu.Kelvin
# Date:         2020/4/7
#-------------------------------------------------------------------------------

import torch
import torch.nn as nn
import numpy as np
import pickle
import torch.nn.functional as F
import os
class SimpleCNN(nn.Module):
    def __init__(self, weight_path='simple_cnn.weights', eps_cnn=1e-5, momentum_cnn=0.05):
        super(SimpleCNN, self).__init__()
        # init and load pre-trained model
        weights = self.load_weight(weight_path)
        self.conv1 = self.init_conv(1, 64, weights['conv1'], weights['b1'])
        self.conv1_bn = nn.BatchNorm2d(num_features=64, eps=eps_cnn, affine=True, momentum=momentum_cnn)
        self.conv2 = self.init_conv(64, 64, weights['conv2'], weights['b2'])
        self.conv2_bn = nn.BatchNorm2d(num_features=64, eps=eps_cnn, affine=True, momentum=momentum_cnn)
        self.conv3 = self.init_conv(64, 64, weights['conv3'], weights['b3'])
        self.conv3_bn = nn.BatchNorm2d(num_features=64, eps=eps_cnn, affine=True, momentum=momentum_cnn)
        self.conv4 = self.init_conv(64, 64, weights['conv4'], weights['b4'])
        self.conv4_bn = nn.BatchNorm2d(num_features=64, eps=eps_cnn, affine=True, momentum=momentum_cnn)

    def load_weight(self, path):
        return pickle.load(open(path, 'rb'))

    def forward(self, X):
        out = F.relu(self.conv1(X))
        out = self.conv1_bn(out)
        out = F.relu(self.conv2(out))
        out = self.conv2_bn(out)
        out = F.relu(self.conv3(out))
        out = self.conv3_bn(out)
        out = F.relu(self.conv4(out))
        out = self.conv4_bn(out)
        out = out.view(-1, 64, 36)

        return torch.mean(out, 2)

    def convert_to_torch_weight(self, weight):
        return np.transpose(weight, [3, 2, 0, 1])

    def init_conv(self, inp, out, weight, bias, convert=True):
        conv = nn.Conv2d(inp, out, 3, 2, 1, bias=True)
        #if convert:
         #   weight = self.convert_to_torch_weight(weight)
        #conv.weight.data = torch.Tensor(weight).float()
        #conv.bias.data = torch.Tensor(bias).float()
        return conv

if __name__ == '__main__':
    data = os.path.dirname(__file__) + '/data'
    simple_cnn = SimpleCNN(weight_path=os.path.join(data,'pretrained_maml.weights'), eps_cnn=1e-5, momentum_cnn=0.05)
    npo = np.random.random((3, 1, 84, 84))
    x = torch.tensor(npo, dtype=torch.float32).float()
    y=simple_cnn(x)
    print(y.shape)