Getting model weights for a particular layer is straightforward. All the model weights can be accessed through the state_dict function. The state_dict function returns a dictionary, with keys as its layers and weights as its values. The following code demonstrates how to pull weights for a particular layer and visualize them:
vgg.state_dict().keys()
cnn_weights = vgg.state_dict()['features.0.weight'].cpu()
The preceding code provides us with the following output:
Each box represents weights of a filter that is of size 3 x 3. Each filter is trained to identify certain patterns in the images.