I am using an autoencoder,
def __init__(self):
super(autoencoder, self).__init__()
#self.channel = channel
self.encoder = nn.Sequential(
nn.Conv2d(in_channels = 3, out_channels = 32, kernel_size = 3, stride = 1), # 26
nn.ReLU(),
nn.MaxPool2d(2, stride = 2), # 13
nn.Conv2d(in_channels = 32, out_channels = 16, kernel_size = 3, stride = 1), # 11
nn.ReLU(),
nn.MaxPool2d(2, stride = 2), # 5
nn.Conv2d(in_channels = 16, out_channels = 8, kernel_size = 3, stride = 1), # 3
nn.ReLU(),
) # 8x3x3
self.decoder = nn.Sequential(
nn.ConvTranspose2d(in_channels = 8, out_channels = 16, kernel_size = 3, stride = 1), #5
nn.ReLU(),
nn.ConvTranspose2d(in_channels = 16, out_channels = 32, kernel_size = 3, stride = 1), #7
nn.ReLU(),
nn.ConvTranspose2d(in_channels = 32, out_channels = 32, kernel_size = 3, stride = 2), #15
nn.ReLU(),
nn.ConvTranspose2d(in_channels= 32, out_channels = 3, kernel_size = 3, stride = 2,padding=2, output_padding=1), #28
nn.Tanh()
)
self.classifier = nn.Sequential(
# nn.AdaptiveAvgPool2d((1,1)),
nn.Flatten(),
nn.Linear(8*3*3,8),
nn.ReLU(),
nn.Linear(8,10)
)
def forward(self, x):
x= self.encoder(x)
out = self.classifier(x)
# feature = self.flatten(x)
x = self.decoder(x)
return x,out
now i want to update the weights of the autoencoder without the classifier using adam optimizer and calculating the L2 loss
between the reconstruction images and the original image.
and want to train only the classifier using adam only, but using different loss i.e, crossentropyloss
. now how can I train encoder, decoder, and the classifier separately in one go?
I want to do something like this:
for epoch in range(10):
print(f"Epoch {epoch+1}")
for images in tqdm(trainloader):
optimizer1.zero_grad()
optimizer2.zero_grad()
images,Ys = images
images = images.cuda()
Ys = Ys.cuda()
x, out = model(images)
loss = loss_fn1(x, images) # l2 norm
loss1 = loss_fn2(out, Ys) # cross entropy
total_loss+=loss
total_loss1+=loss1
# print('Loss b/w images: ',loss.item())
# print('Loss b/w labels: ',loss1.item())
# break
loss.backward() # l2 norm
loss1.backward() # cross en
optimizer1.step()
optimizer2.step()
But it is showing error, which I fixed by setting retrain_graph=True
, but that was not feasible since it is taking all my GPU memory. Is there any other alternative?