https://github.com/Wulingtian/EfficientNetv2_TensorRT_int8/blob/master/effnetv2.py
测试代码:
if __name__ == '__main__':
import os
import time
model=EfficientNet(5)
model.cuda()
model.eval()
model_path = "dicenet.pth"
torch.save(model.state_dict(), model_path)
fsize = os.path.getsize(model_path)
fsize = fsize / float(1024 * 1024)
print(f"model size {round(fsize, 2)} m")
size=224
input = torch.rand(10, 3, size, size).cuda()
for i in range(15):
t1 = time.time()
loc = model(input)
cnt = time.time() - t1
print(cnt, loc.size())
</