-
Notifications
You must be signed in to change notification settings - Fork 0
/
encoder.py
27 lines (24 loc) · 846 Bytes
/
encoder.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
import numpy as np
import cv2
import matplotlib.pyplot as plt
import os
import torch
from torchvision import models
import torch.nn as nn
import torchvision.models as models
from torchvision import transforms
import PIL
class EncoderCNN(nn.Module):
def __init__(self, embed_size):
super(EncoderCNN, self).__init__()
resnet = models.resnet50(pretrained=True)
for param in resnet.parameters():
param.requires_grad_(False)
modules = list(resnet.children())[:-1]
self.resnet = nn.Sequential(*modules)
# self.embed = nn.Linear(resnet.fc.in_features, embed_size)
def forward(self, images):
features = self.resnet(images)
features = features.view(features.size(0), -1)
# features = self.embed(features)
return features