Pass an arbitrary image size to cnn in pytorch - python

I'm trying to train a lenet model in pytorch, The ideia is to put images of any size in it, so I started doing with nn.AdaptiveAvgPool2d but the error comes as
mat1 dim 1 must match mat2 dim 0
Here is my code
class LeNet5(nn.Module):
def __init__(self, num_classes=10):
super(LeNet5, self).__init__()
self.conv_1 = nn.Conv2d(
in_channels=1, out_channels=32, kernel_size=5, bias=False
)
self.relu_1 = nn.ReLU(inplace=True)
self.maxpool_1 = nn.MaxPool2d(kernel_size=2, stride=2)
self.conv_2 = nn.Conv2d(
in_channels=32, out_channels=256, kernel_size=5, bias=False
)
self.relu_2 = nn.ReLU(inplace=True)
self.maxpool_2 = nn.MaxPool2d(kernel_size=2, stride=2)
self.avgpool = nn.AdaptiveAvgPool2d(output_size=1)
self.flatten = nn.Flatten()
self.fc_1 = nn.Linear(in_features=4096, out_features=120, bias=False)
self.fc_2 = nn.Linear(in_features=120, out_features=84)
self.fc_3 = nn.Linear(in_features=84, out_features=num_classes)
def forward(self, input):
conv_1_output = self.conv_1(input)
relu_1_output = self.relu_1(conv_1_output)
maxpool_1_output = self.maxpool_1(relu_1_output)
conv_2_output = self.conv_2(maxpool_1_output)
relu_2_output = self.relu_2(conv_2_output)
maxpool_2_output = self.maxpool_2(relu_2_output)
flatten_output = self.flatten((self.avgpool(maxpool_2_output).view(maxpool_2_output.shape[0], -1)))
fc_1_output = self.fc_1(flatten_output)
fc_2_output = self.fc_2(fc_1_output)
fc_3_output = self.fc_3(fc_2_output)
return fc_3_output

if you read the theory on AdaptiveAvgPool2d, this is what it says " we specify the output size And the stride and kernel-size are automatically selected to adapt to the needs"
More info available here
Hence Your spatial dimension is reduced by AdaptiveAvgPool2d and not the depth of feature maps.
So, the spatial dimension will be 1x1 and depth will still be 256 , making your
self.fc_1 = nn.Linear(in_features=256, out_features=120, bias=False) and not self.fc_1 = nn.Linear(in_features=4096, out_features=120, bias=False)

Related

Trying to understand this CNN Model - PyTorch

I am trying to understand this model, by building it from scratch.
I am trying to start with just one layer, but I am getting error stating output of layer1 is not the expected input type for fc1.
RuntimeError: mat1 and mat2 shapes cannot be multiplied (64x197136 and 2704x2704)
I am not 100% sure if I can do this with 1 layer, but I am trying to learn by playing around with it.
code(Which I am trying to play around with):
# Craete a neural network from pytorch
# https://www.kaggle.com/code/reukki/pytorch-cnn-tutorial-with-cats-and-dogs
class Cnn(nn.Module):
def __init__(self):
super(Cnn,self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, padding=0, stride=2),
nn.BatchNorm2d(num_features=16),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2)
)
self.fc1 = nn.Linear(3*3*16,10)
self.dropout = nn.Dropout(0.5)
self.fc2 = nn.Linear(10,2)
self.relu = nn.ReLU()
def forward(self,x):
out = self.layer1(x)
out = out.view(out.size(0),-1)
out = self.fc1(out)
out = self.relu(out)
out = self.fc2(out)
return out
I am also unsure what out.view(out.size(0),-1) does.
Would be great if anyone can help me understand better. Looking forward to some
[Update]
Kaggle Notebook link:- https://www.kaggle.com/austonpramodh/cats-dogs-cnn-learning
out.view(out.size(0),-1) flattens your output dimension to 1-dimension(1d) according to (whatever_number_fits_it_perfectly_for_1d), for example if out.size(0) is 512x512, this command will make it to 262144 and remove the 2nd dimension. You can also say it as 262144x1.
I was able to get the answer from one of my friend, Thanks to Vihari.
class Cnn(nn.Module):
def __init__(self):
super(Cnn,self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, padding=0, stride=2),
nn.BatchNorm2d(num_features=16),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2)
)
self.fc1 = nn.Linear(55 * 55 * 16, 2)
self.relu = nn.ReLU()
def forward(self,x):
out = self.layer1(x)
# print("****Layer Size***")
# # 55*55*16
# print(out.size(1))
# print("****Layer Size***")
out = self.relu(out)
out = out.view(out.size(0), out.size(1)*out.size(2)* out.size(3))
# print("****View Size***")
# print(out.size()[-1], out.size(1))
# print("****View Size***")
out = self.fc1(out)
return out
Also, Thanks for the out.view explanation #Pranav.

CycleGan only producing black images

I'm trying to follow a Cyclegan tutorial here:https://www.youtube.com/watch?v=4LktBHGCNfw. It also has the Github to his code. At first I thought I had a typo, but then I tried copying and pasting his code and still the results are only black images. Is anybody else able to copy the code and get actual images?
Generator model
class ConvBlock(nn.Module):
def __init__(self, in_channels, out_channels, down=True, use_act=True, **kwargs):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, padding_mode="reflect", **kwargs)
if down
else nn.ConvTranspose2d(in_channels, out_channels, **kwargs),
nn.InstanceNorm2d(out_channels),
nn.ReLU(inplace=True) if use_act else nn.Identity()
)
def forward(self, x):
return self.conv(x)
class ResidualBlock(nn.Module):
def __init__(self, channels):
super().__init__()
self.block = nn.Sequential(
ConvBlock(channels, channels, kernel_size=3, padding=1),
ConvBlock(channels, channels, use_act=False, kernel_size=3, padding=1),
)
def forward(self, x):
return x + self.block(x)
class Generator(nn.Module):
def __init__(self, img_channels, num_features = 64, num_residuals=9):
super().__init__()
self.initial = nn.Sequential(
nn.Conv2d(img_channels, num_features, kernel_size=7, stride=1, padding=3, padding_mode="reflect"),
nn.InstanceNorm2d(num_features),
nn.ReLU(inplace=True),
)
self.down_blocks = nn.ModuleList(
[
ConvBlock(num_features, num_features*2, kernel_size=3, stride=2, padding=1),
ConvBlock(num_features*2, num_features*4, kernel_size=3, stride=2, padding=1),
]
)
self.res_blocks = nn.Sequential(
*[ResidualBlock(num_features*4) for _ in range(num_residuals)]
)
self.up_blocks = nn.ModuleList(
[
ConvBlock(num_features*4, num_features*2, down=False, kernel_size=3, stride=2, padding=1, output_padding=1),
ConvBlock(num_features*2, num_features*1, down=False, kernel_size=3, stride=2, padding=1, output_padding=1),
]
)
self.last = nn.Conv2d(num_features*1, img_channels, kernel_size=7, stride=1, padding=3, padding_mode="reflect")
def forward(self, x):
x = self.initial(x)
for layer in self.down_blocks:
x = layer(x)
x = self.res_blocks(x)
for layer in self.up_blocks:
x = layer(x)
return torch.tanh(self.last(x))
Discriminator Model
class Block(nn.Module):
def __init__(self, in_channels, out_channels, stride):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, 4, stride, 1, bias=True, padding_mode="reflect"),
nn.InstanceNorm2d(out_channels),
nn.LeakyReLU(0.2, inplace=True),
)
def forward(self, x):
return self.conv(x)
class Discriminator(nn.Module):
def __init__(self, in_channels=3, features=[64, 128, 256, 512]):
super().__init__()
self.initial = nn.Sequential(
nn.Conv2d(
in_channels,
features[0],
kernel_size=4,
stride=2,
padding=1,
padding_mode="reflect",
),
nn.LeakyReLU(0.2, inplace=True),
)
layers = []
in_channels = features[0]
for feature in features[1:]:
layers.append(Block(in_channels, feature, stride=1 if feature==features[-1] else 2))
in_channels = feature
layers.append(nn.Conv2d(in_channels, 1, kernel_size=4, stride=1, padding=1, padding_mode="reflect"))
self.model = nn.Sequential(*layers)
def forward(self, x):
x = self.initial(x)
return torch.sigmoid(self.model(x))
Train:
def train_fn(disc_H, disc_Z, gen_Z, gen_H, loader, opt_disc, opt_gen, l1, mse, d_scaler, g_scaler):
H_reals = 0
H_fakes = 0
loop = tqdm(loader, leave=True)
for idx, (zebra, horse) in enumerate(loop):
zebra = zebra.to(config.DEVICE)
horse = horse.to(config.DEVICE)
# Train Discriminators H and Z
with torch.cuda.amp.autocast():
fake_horse = gen_H(zebra)
D_H_real = disc_H(horse)
D_H_fake = disc_H(fake_horse.detach())
H_reals += D_H_real.mean().item()
H_fakes += D_H_fake.mean().item()
D_H_real_loss = mse(D_H_real, torch.ones_like(D_H_real))
D_H_fake_loss = mse(D_H_fake, torch.zeros_like(D_H_fake))
D_H_loss = D_H_real_loss + D_H_fake_loss
fake_zebra = gen_Z(horse)
D_Z_real = disc_Z(zebra)
D_Z_fake = disc_Z(fake_zebra.detach())
D_Z_real_loss = mse(D_Z_real, torch.ones_like(D_Z_real))
D_Z_fake_loss = mse(D_Z_fake, torch.zeros_like(D_Z_fake))
D_Z_loss = D_Z_real_loss + D_Z_fake_loss
# put it togethor
D_loss = (D_H_loss + D_Z_loss)/2
opt_disc.zero_grad()
d_scaler.scale(D_loss).backward()
d_scaler.step(opt_disc)
d_scaler.update()
# Train Generators H and Z
with torch.cuda.amp.autocast():
# adversarial loss for both generators
D_H_fake = disc_H(fake_horse)
D_Z_fake = disc_Z(fake_zebra)
loss_G_H = mse(D_H_fake, torch.ones_like(D_H_fake))
loss_G_Z = mse(D_Z_fake, torch.ones_like(D_Z_fake))
# cycle loss
cycle_zebra = gen_Z(fake_horse)
cycle_horse = gen_H(fake_zebra)
cycle_zebra_loss = l1(zebra, cycle_zebra)
cycle_horse_loss = l1(horse, cycle_horse)
# identity loss (remove these for efficiency if you set lambda_identity=0)
identity_zebra = gen_Z(zebra)
identity_horse = gen_H(horse)
identity_zebra_loss = l1(zebra, identity_zebra)
identity_horse_loss = l1(horse, identity_horse)
# add all togethor
G_loss = (
loss_G_Z
+ loss_G_H
+ cycle_zebra_loss * config.LAMBDA_CYCLE
+ cycle_horse_loss * config.LAMBDA_CYCLE
+ identity_horse_loss * config.LAMBDA_IDENTITY
+ identity_zebra_loss * config.LAMBDA_IDENTITY
)
opt_gen.zero_grad()
g_scaler.scale(G_loss).backward()
g_scaler.step(opt_gen)
g_scaler.update()
if idx % 200 == 0:
save_image(fake_horse*0.5+0.5, f"saved_images/horse_{idx}.png")
save_image(fake_zebra*0.5+0.5, f"saved_images/zebra_{idx}.png")
#loop.set_postfix(H_real=H_reals/(idx+1), H_fake=H_fakes/(idx+1))
def main():
disc_H = Discriminator(in_channels=3).to(config.DEVICE)
disc_Z = Discriminator(in_channels=3).to(config.DEVICE)
gen_Z = Generator(img_channels=3, num_residuals=9).to(config.DEVICE)
gen_H = Generator(img_channels=3, num_residuals=9).to(config.DEVICE)
opt_disc = optim.Adam(
list(disc_H.parameters()) + list(disc_Z.parameters()),
lr=config.LEARNING_RATE,
betas=(0.5, 0.999),
)
opt_gen = optim.Adam(
list(gen_Z.parameters()) + list(gen_H.parameters()),
lr=config.LEARNING_RATE,
betas=(0.5, 0.999),
)
L1 = nn.L1Loss()
mse = nn.MSELoss()
if config.LOAD_MODEL:
load_checkpoint(
config.CHECKPOINT_GEN_H, gen_H, opt_gen, config.LEARNING_RATE,
)
load_checkpoint(
config.CHECKPOINT_GEN_Z, gen_Z, opt_gen, config.LEARNING_RATE,
)
load_checkpoint(
config.CHECKPOINT_CRITIC_H, disc_H, opt_disc, config.LEARNING_RATE,
)
load_checkpoint(
config.CHECKPOINT_CRITIC_Z, disc_Z, opt_disc, config.LEARNING_RATE,
)
dataset = HorseZebraDataset(
root_horse=config.TRAIN_DIR+"/horses", root_zebra=config.TRAIN_DIR+"/zebras", transform=config.transforms
)
#val_dataset = HorseZebraDataset(
#root_horse="cyclegan_test/horse1", root_zebra="cyclegan_test/zebra1", transform=config.transforms
#)
#val_loader = DataLoader(
#val_dataset,
#batch_size=1,
#shuffle=False,
#pin_memory=True,
#)
loader = DataLoader(
dataset,
batch_size=config.BATCH_SIZE,
shuffle=True,
#num_workers=config.NUM_WORKERS,
pin_memory=True
)
g_scaler = torch.cuda.amp.GradScaler()
d_scaler = torch.cuda.amp.GradScaler()
for epoch in range(config.NUM_EPOCHS):
train_fn(disc_H, disc_Z, gen_Z, gen_H, loader, opt_disc, opt_gen, L1, mse, d_scaler, g_scaler)
if config.SAVE_MODEL:
save_checkpoint(gen_H, opt_gen, filename=config.CHECKPOINT_GEN_H)
save_checkpoint(gen_Z, opt_gen, filename=config.CHECKPOINT_GEN_Z)
save_checkpoint(disc_H, opt_disc, filename=config.CHECKPOINT_CRITIC_H)
save_checkpoint(disc_Z, opt_disc, filename=config.CHECKPOINT_CRITIC_Z)

Autoencoder outputs B&W images from color

We are in the process of training an AE on CIFAR10 images. We used the following architecture:
class OurAE(nn.Module):
def __init__(self, in_channels, z_channels):
super(OurAE, self).__init__()
self.tot_diff = None
self.in_channels = in_channels
curr_channels = in_channels
encoder = []
channels = [3, 16, 64] + [z_channels]
for out_channels in channels:
encoder += [
nn.Conv2d(in_channels=curr_channels, out_channels=out_channels, kernel_size=3, padding=1, stride=2),
nn.ReLU()
]
curr_channels = out_channels
self.encoder = nn.Sequential(*encoder)
curr_channels = z_channels
decoder = []
channels = [64, 16, 3] + [in_channels]
for out_channels in channels:
decoder += [
nn.ConvTranspose2d(in_channels=curr_channels, out_channels=out_channels, kernel_size=4, padding=1, stride=2),
nn.ReLU()
]
curr_channels = out_channels
decoder = decoder[:-1] # removing the RELU layer
decoder.append(nn.Sigmoid())
self.decoder = nn.Sequential(*decoder)
def forward(self, x):
return self.decoder(self.encoder(x))
We are not sure why but we always get black and white images.
We tried to replace the Sigmoid with ReLU in the last layer but to no avail.
These are the loss function and the optimizers that we used:
optimizer = torch.optim.Adam(classifier.parameters(), lr=lr)
criterion = torch.nn.CrossEntropyLoss()
Here is an example of an input and output of the AE after training:
I had the same problem. I changed the loss function to 'mae'. That cleared the issue for me.

Pytorch: RuntimeError: Sizes of tensors must match except in dimension 2

I´ve got a problem when passing an image to my unet. I get the following error :
Traceback (most recent call last):
File "path\Main.py", line 101, in <module>
outputs = model(inputs[None,...].float())
File "C:\Users\...\AppData\Local\Programs\Python\Python39\lib\site-packages\torch\nn\modules\module.py", line 727, in _call_impl
result = self.forward(*input, **kwargs)
File "path\UNets.py", line 53, in forward
upconv2 = self.upconv2(torch.cat([upconv1,conv4]),1)
RuntimeError: Sizes of tensors must match except in dimension 2. Got 15 and 10 (The offending index is 0)
The images i use are of different sizes, i format them to 256x256. They are grayscale images
my dataloading:
def getImageAndTransform(self,item):
data = Image.open(self.datalist[item])
label = Image.open(self.labellist[item])
p = torchvision.transforms.Compose([torchvision.transforms.Scale((256,256))])
data = p(data)
label = p(label)
data = torch.from_numpy(np.array(data))
label = torch.from_numpy(np.array((label)))
return data, label
my Unet:
class Unet(SegmentationNetwork):
def __init__(self,config):
super(StandardUnet,self).__init__(config = config)
#down
self.downconv1 =self.contract_block(self.in_channels,self.channels[0],self.kernel[0],self.padding[0])
self.downconv2 =self.contract_block(self.channels[0],self.channels[1],self.kernel[1],self.padding[1])
self.downconv3 =self.contract_block(self.channels[1],self.channels[2],self.kernel[2],self.padding[2])
self.downconv4 =self.contract_block(self.channels[2],self.channels[3],self.kernel[3],self.padding[3])
self.downconv5 =self.contract_block(self.channels[3],self.channels[4],self.kernel[4],self.padding[4])
#up
self.upconv1 = self.expand_block(self.channels[4],self.channels[3],self.kernel[4],self.padding[4])
self.upconv2 = self.expand_block(self.channels[3],self.channels[2],self.kernel[3],self.padding[3])
self.upconv3 = self.expand_block(self.channels[2], self.channels[1], self.kernel[2], self.padding[2])
self.upconv4 = self.expand_block(self.channels[1], self.channels[0], self.kernel[1], self.padding[1])
self.upconv5 = self.expand_block(self.channels[0], self.out_channels, self.kernel[0], self.padding[0])
def forward(self,x):
#down
conv1 = self.downconv1(x)
conv2 = self.downconv2(conv1)
conv3 = self.downconv3(conv2)
conv4 = self.downconv4(conv3)
conv5 = self.downconv5(conv4)
#up
upconv1 = self.upconv1(conv5)
upconv2 = self.upconv2(torch.cat([upconv1,conv4]),1)
upconv3 = self.upconv3(torch.cat([upconv2,conv3]),1)
upconv4 = self.upconv4(torch.cat([upconv3,conv2]),1)
upconv5 = self.upconv5(torch.cat([upconv4,conv1]),1)
self.out = upconv5
def contract_block(self,in_channels,out_channels,kernel_size, padding):
contract = nn.Sequential(
nn.Conv2d(in_channels,out_channels,kernel_size=kernel_size,stride=1,padding=padding),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels,out_channels,kernel_size=kernel_size,stride=1,padding=padding),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3,stride=2,padding=1))
return contract
def expand_block(self,in_channels,out_channels,kernel_size,padding):
expand = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=padding),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size, stride=1, padding=padding),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
nn.ConvTranspose2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1, output_padding=1)
)
return expand
my implementation:
for i, data in enumerate(dataloader_train, 0): # inputdata as list of [inputs,labels]
data[0].size()
data[1].size()
inputs, labels = data[0].to(device), data[1].to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
outputs = model(inputs[None,...].float())
loss = loss_function(outputs, labels)
loss.backward()
optimizer.step()
Can anyone tell me what i should do to fix this?
Maybe the question is stupid, but since i´m new to torch and deep learning in general i would apprechiate help
Thanks
I found my error, there was a bracket in the wrong place in the upconvolution step.
correct it would be upconv2 = self.upconv2(torch.cat([upconv1,conv4],1))

Unknown behaviour of HOOKS in PyTorch

I have a straightforward and simple CNN below,
# creat a dummy deep net
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1,2, kernel_size=3, stride=1, padding=1, bias=True)
self.conv2 = nn.Conv2d(2,3, kernel_size=3, stride=1, padding=1, bias=True)
self.conv3 = nn.Conv2d(3,1, kernel_size=3, stride=1, padding=1, bias=True)
self.seq = nn.Sequential(
nn.Conv2d(1,5, kernel_size=3, stride=1, padding=1, bias=True),
nn.LeakyReLU(negative_slope=0.2, inplace=True),
nn.Conv2d(5,1, kernel_size=3, stride=1, padding=1, bias=True),
)
self.relu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
def forward(self, x):
out = self.relu(self.conv1(x))
out = self.conv3(self.conv2(out))
out = out + x
out = self.seq(x)
return out
5 hooks have been applied to each layer for the forward pass.
Hooked 0 to Conv2d(1, 2, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
Hooked 1 to Conv2d(2, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
Hooked 2 to Conv2d(3, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
Hooked 3 to Sequential(
(0): Conv2d(1, 5, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): LeakyReLU(negative_slope=0.2, inplace=True)
(2): Conv2d(5, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
)
Hooked 4 to LeakyReLU(negative_slope=0.2, inplace=True)
These hooks have been created using following class
# ------------------The Hook class begins to calculate each layer stats
class Hook():
def __init__(self, module, backward=False):
if backward==False:
self.hook = module.register_forward_hook(self.hook_fn)
else:
self.hook = module.register_backward_hook(self.hook_fn)
self.inputMean = []
self.outputMean = []
def hook_fn(self, module, input, output):
self.inputMean.append(input[0][0,...].mean().item())#calculate only for 1st image in the batch
print('\nIn hook class input {}'.format(input[0].size()))
self.outputMean.append(output[0][0,...].mean().item())
print('In hook class outout {}'.format(output[0].size()))
# create hooks on each layer
hookF = []
for i,layer in enumerate(list(net.children())):
print('Hooked to {}'.format(layer))
hookF.append(Hook(layer))
Please note between Hook 1 and Hook 2 there is no ReLU
self.conv3(self.conv2(out)). Thus OUTPUT of HOOK1 is INPUT to HOOK2 and should be identical. BUT THIS DOES NOT TURNS OUT TO BE WHY? Below is output for HOOK1 and HOOK2
Hook of layer 1 (HOOK on layer 1 which is self.conv2)
... OutputMean: [0.2381615787744522, 0.2710852324962616, 0.30706286430358887, 0.26064932346343994, 0.24395985901355743]
Hook of layer 2 (HOOK on layer 2 which is self.conv3)
InputMean: [0.13127394020557404, 0.1611362248659134, 0.1457807868719101, 0.17380955815315247, 0.1537724733352661], OutputMean: ...
These two values should have been the same but do not turn out to be.
------ The Full code is shown below -------
import torch
import torch.nn as nn
# creat a dummy deep net
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1,2, kernel_size=3, stride=1, padding=1, bias=True)
self.conv2 = nn.Conv2d(2,3, kernel_size=3, stride=1, padding=1, bias=True)
self.conv3 = nn.Conv2d(3,1, kernel_size=3, stride=1, padding=1, bias=True)
self.seq = nn.Sequential(
nn.Conv2d(1,5, kernel_size=3, stride=1, padding=1, bias=True),
nn.LeakyReLU(negative_slope=0.2, inplace=True),
nn.Conv2d(5,1, kernel_size=3, stride=1, padding=1, bias=True),
)
self.relu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
def forward(self, x):
out = self.relu(self.conv1(x))
out = self.conv3(self.conv2(out))
out = out + x
out = self.seq(x)
return out
net = Net()
print(net)
criterion = nn.MSELoss()
# ------------------The Hook class begins to calculate each layer stats
class Hook():
def __init__(self, module, backward=False):
if backward==False:
self.hook = module.register_forward_hook(self.hook_fn)
else:
self.hook = module.register_backward_hook(self.hook_fn)
self.inputMean = []
self.outputMean = []
def hook_fn(self, module, input, output):
self.inputMean.append(input[0][0,...].mean().item())#calculate only for 1st image in the batch
print('\nIn hook class input {}'.format(input[0].size()))
self.outputMean.append(output[0][0,...].mean().item())
print('In hook class outout {}'.format(output[0].size()))
# create hooks on each layer
hookF = []
for i,layer in enumerate(list(net.children())):
print('Hooked to {}'.format(layer))
hookF.append(Hook(layer))
optimizer = torch.optim.Adam(net.parameters())
# Do 5 forward pass
for _ in range(5):
print('Iteration --------')
data = torch.rand(2,1,10,10)*10
print('Input mean is {}'.format(data[0,...].mean()))
target = data.clone()
out = net(data)
loss = criterion(out, target)
print('backward')
loss.backward()
optimizer.step()
optimizer.zero_grad()
for i,h in enumerate(hookF):
print('\n Hook of layer {}'.format(i))
print('InputMean: {}, OutputMean: {}'.format(h.inputMean, h.outputMean))
h.hook.remove()
The problem is that in your Conv2d layer input is a tuple and output is a torch.Tensor. Therefore output[0][0,...] is selecting the first item from dim 0 in the tensor whereas input[0][0,...] is selecting the first item from the tuple.
You just need to change output[0][0,...] to output[0,...].

Categories