def __init__(self, args, loader, model, ckp):
self.args = args
self.ckp = ckp
self.model = model
self.my_model = self.model.get_model()
self.my_model.reset()
self.loader_train, self.loader_test = loader
self.device = torch.device('cpu' if args.cpu_only else 'cuda')
self.loss = nn.CrossEntropyLoss()
self.encoder = EncoderCNN.make_model(args).to(self.device)
self.decoder = self.my_model.to(self.device)
# IF gradient only in RNN + linear layer:
if not self.args.fine_tune:
print("Not fine tuning CNN")
self.params = list(self.decoder.parameters()) + list(
self.encoder.linear.parameters()) + list(
self.encoder.bn.parameters())
else:
print("Fine tuning CNN")
self.params = list(self.decoder.parameters()) + list(
self.encoder.parameters())
# only train decoder
self.optimizer = optim.Adam(self.params, args.learning_rate) # ADAM
# self.optimizer = optim.SGD(self.params, args.learning_rate, momentum=0.8) # SGD
self.lr_scheduler = set_scheduler(args, self.optimizer)
if args.cpu_only:
kwargs = {'map_location': lambda storage, loc: storage}
else:
kwargs = {}
if args.load_path != '.':
if args.pre_train != '.':
print('Load model from : {}'.format(args.pre_train))
self.my_model.load_state_dict(torch.load(
args.pre_train, **kwargs),
strict=False)
else:
if not self.args.test_only:
print('Load model from : {}'.format(
os.path.join(ckp.log_dir, 'model', 'model_latest.pt')))
self.my_model.load_state_dict(torch.load(
os.path.join(ckp.log_dir, 'model', 'model_latest.pt'),
**kwargs),
strict=False)
else:
print('Load model from : {}'.format(
os.path.join(ckp.log_dir, 'model', 'model_best.pt')))
self.my_model.load_state_dict(torch.load(
os.path.join(ckp.log_dir, 'model', 'model_best.pt'),
**kwargs),
strict=False)
if not self.args.test_only:
self.optimizer.load_state_dict(
torch.load(os.path.join(ckp.log_dir, 'optimizer.pt')))
for _ in range(1, len(ckp.loss.result)):
self.lr_scheduler.step()
import torch
import numpy as np
import cv2
from model import EncoderCNN
from model import CRNNv2
from option import args
from data import common
device = torch.device('cpu')
encoder = EncoderCNN.make_model(args)
encoder.eval()
decoder = CRNNv2.make_model(args)
decoder.eval()
encoder.load_state_dict(torch.load('./Inference_model/model_enc_best.pt', map_location='cpu'))
decoder.load_state_dict(torch.load('./Inference_model/model_best.pt', map_location='cpu'))
print("Model Loaded")
#read image
img = np.zeros((1,96,480))
img_read = cv2.imread('trial.png',0)
img_read = cv2.resize(img_read, (480, 96), interpolation = cv2.INTER_NEAREST)
img_read[img_read>100]=255
cv2.imwrite('input_img.png',img_read)
img_read = common.normalize_img(img_read)
print(img_read.shape)
img_read_expanded = np.expand_dims(img_read, axis=0)
img = img_read_expanded
img = np.expand_dims(img, axis=0)
