def prepare_fpn_model(given_image):
'''
准备代理模型
:param given_image: layout: x2y
:return: model and device
'''
import sys
sys.path.append("..")
import os
from fpn.model import fpn
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
device = torch.device(
'cuda:0') if torch.cuda.is_available() else torch.device('cpu')
model = fpn().to(device)
if given_image == 'layout':
model_path = 'modelfile/fpn_x2y_ex.pth'
# model_path = 'modelfile/fpn_x2y_5w.pth'
print("model path:", model_path)
if torch.cuda.is_available():
model.load_state_dict(torch.load(model_path, map_location='cuda:0'))
else:
model.load_state_dict(torch.load(model_path, map_location='cpu'))
model.eval()
return model, device
return mae, accuracy
if __name__ == "__main__":
import os
from fpn.model import fpn
from utils import project_path
from utils.mat2pic import TestDataset, GeneralDataset, trans_separate
from torch.utils.data import DataLoader
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
device = torch.device('cuda')
if not torch.cuda.is_available():
print("Use CPU")
device = torch.device('cpu')
model = fpn().to(device)
model_path = os.path.join(project_path, 'data', 'fpn.pth.52')
dataset_test = TestDataset(trans_separate, resize_shape=(200, 200))
valid_loader = DataLoader(dataset_test,
batch_size=1,
shuffle=False,
drop_last=True)
print("model path:", model_path)
if torch.cuda.is_available():
model.load_state_dict(torch.load(model_path, map_location='cuda:0'))
else:
model.load_state_dict(torch.load(model_path, map_location='cpu'))