Skip to content

Instantly share code, notes, and snippets.

@AFutureD

AFutureD/pytorch_model_shape_test.py

Last active January 16, 2020 03:35
Show Gist options
  • Save AFutureD/442b4b8457f553ece2fca5ddeedbf5f4 to your computer and use it in GitHub Desktop.
Save AFutureD/442b4b8457f553ece2fca5ddeedbf5f4 to your computer and use it in GitHub Desktop.
Download ZIP
Pytorch test model's shape
Raw
pytorch_model_shape_test.py
img = np.random.randint(0,256,(4,3,512,512)) # (4,3,512,512) => batch size, channel, height, width
img = torch.from_numpy(img).float().div(255) # convert numpy to tensor and scale to 0-1.
cbam_tmp = resnet50_cbam() # create a model.
cbam_tmp(img)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment