atom-predict/msunet/.ipynb_checkpoints/test_unet-checkpoint.py

36 lines
642 B
Python
Executable File

import glob
import time
import timm
import torch
import collections
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from torchinfo import summary as ts
from thop import profile
from sklearn.manifold import TSNE
from core.model import *
from core.data import *
from core.metrics import *
from torch_geometric.loader import DataLoader
from torch_geometric.nn import summary as nns
import torch.nn as nn
model = C_FCRN_Aux(3).cuda().eval()
x = torch.randn(128, 3, 256, 256).cuda()
nums = 100
s = time.time()
for i in range(nums):
with torch.no_grad():
model(x)
e = time.time()
print((e - s) / nums)