machrisaa / tensorflow-vgg

VGG19 and VGG16 on Tensorflow
2.21k stars 1.08k forks source link

npy file for pytorch? #76

Closed HeConnor closed 3 years ago

HeConnor commented 3 years ago

Hello sir,

I rewritten the VGG19.py using pytorch, and used the model weight from VGG19.npy.

But when I test my vgg19, the results were error, different with tensorflow,....

I checked my code carefully, but found nothing, are there any difference in operator (Conv2d, max_pool, ...) between tensorflow and pytorch? (I had seen the documention of tensorflow and pytorch,...)

the weight file(npy) was used as below, could you help me?

# load tensorflow model weights
data_dict = np.load('vgg19.npy', allow_pickle=True, encoding='latin1').item()

# my vgg19 pytorch
model = Vgg19()     
model_dict = model.state_dict()

pretrained_dict = OrderedDict()
for key in model_dict.keys():
    if '.weight' in key:
        data_key = key.replace('.0', '')
        data_key = data_key.replace('.weight', '')
        data = data_dict[data_key][0]
        if len(data.shape) > 2:
            # [h, w, c, b] to [b, c, h, w]
            data = np.ascontiguousarray(data.transpose([3, 2, 0, 1]))
        else:
            data = np.ascontiguousarray(data.transpose([1, 0]))
    elif '.bias' in key:
        data_key = key.replace('.0', '')
        data_key = data_key.replace('.bias', '')
        data = data_dict[data_key][1]
    else:
        raise Exception('Error keys')
    paras = torch.from_numpy(data)
    pretrained_dict[key] = paras

model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)

class Vgg19(nn.Module):
    def __init__(self, model_file=None):
        super(Vgg19, self).__init__()
        self.conv1_1 = self.conv_layer(3, 64, 3)
        self.conv1_2 = self.conv_layer(64, 64, 3)

        self.conv2_1 = self.conv_layer(64, 128, 3)
        self.conv2_2 = self.conv_layer(128, 128, 3)
        ......
        self.max_pool = nn.MaxPool2d(2, 2)

    def conv_layer(self, in_dim, out_dim, kernel_size=3):
        return nn.Sequential(
            nn.Conv2d(in_dim, out_dim, kernel_size, stride=1, padding=1),
            nn.ReLU(inplace=True),
        )
    ......

Best regards!