Hi,
I’m trying to train a simple 2-layer convolutional network to classify Cifar-10 images where layers are trained separately.
For this purpose, I have defined 3 different nn.Modules. Two for the convolutional layers, and one for the classifier at the end:
class Conv1(nn.Module):
def __init__(self):
super(Conv1, self).__init__()
self.conv = nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1)
# self.z1_h = torch.nn.Parameter(Variable(torch.randn(64, 32, 16, 16, device=device)),
# requires_grad=True)
def forward(self, x):
x1 = self.conv(x)
z1 = F.relu(x1)
return x1, z1
class Conv2(nn.Module):
def __init__(self):
super(Conv2, self).__init__()
self.conv = nn.Conv2d(32, 32, kernel_size=3, stride=2, padding=1)
self.z1_h = torch.nn.Parameter(Variable(torch.randn(64, 32, 16, 16, device=device)),
requires_grad=True)
def forward(self, x=None):
if x is not None:
x2 = self.conv(x)
else:
x2 = self.conv(self.z1_h)
z2 = F.relu(x2)
return x2, z2
class Classifier(nn.Module):
def __init__(self):
super(Classifier, self).__init__()
self.fc1 = nn.Linear(32 * 8 * 8, 512)
self.fc2 = nn.Linear(512, 512)
self.fc3 = nn.Linear(512, 10)
self.dropout = nn.Dropout(0.5)
self.z2_h = torch.nn.Parameter(Variable(torch.randn(64, 32, 8, 8, device=device)),
requires_grad=True)
def forward(self, x=None):
# flatten the features
if x is not None:
flat_feat = x.view(-1, 32 * 8 * 8)
else:
flat_feat = self.z2_h.view(-1, 32 * 8 * 8)
fc1 = self.fc1(flat_feat)
fc1 = F.relu(fc1)
fc1 = self.dropout(fc1)
fc2 = self.fc2(fc1)
fc2 = F.relu(fc2)
fc2 = self.dropout(fc2)
out = self.fc3(fc2)
out = F.log_softmax(out, 1)
return out
As it can be seen, I have defined variables self.z1_h and self.z2_h as the input to the 2nd and 3rd layers. These two variables will optimize to be the same as the output of their previous layers, z1 and z2.
To optimize this network, I defined 3 different optimizers to update the parameters in alternation order: W (and bias) of the layers, and the 2 introduced variables self.z1_h and self.z2_h.
optimizer1 = optim.SGD(net1.parameters(), lr=LR, momentum=MOMENTUM, weight_decay=WEIGHT_DECAY)
optimizer2 = optim.SGD(net2.parameters(), lr=LR, momentum=MOMENTUM, weight_decay=WEIGHT_DECAY)
optimizer3 = optim.SGD(net3.parameters(), lr=LR, momentum=MOMENTUM, weight_decay=WEIGHT_DECAY)
# net 1
# zero the parameter gradients
optimizer1.zero_grad()
# forward + backward + optimize
x1, z1 = net1(inputs)
loss1 = F.mse_loss(x1, z1)
loss1.backward(retain_graph=True)
optimizer1.step()
# net 2
# zero the parameter gradients
optimizer2.zero_grad()
# forward + backward + optimize
x2, z2 = net2()
# x2, z2, z2_h = net2(z1_h)
loss2 = F.mse_loss(z1, net2.z1_h) + F.mse_loss(x2, z2)
for param in net2.parameters():
param.requires_grad = True
net2.z1_h.requires_grad = False # z1_h constant
loss2.backward(retain_graph=True)
optimizer2.step()
for param in net2.parameters():
param.requires_grad = False
net2.z1_h.requires_grad = True # W and bias contant
loss2.backward(retain_graph=True)
optimizer2.step()
# net 3
# zero the parameter gradients
optimizer3.zero_grad()
# forward + backward + optimize
outputs = net3()
loss3 = F.nll_loss(outputs, labels) + F.mse_loss(z2, net3.z2_h)
for param in net3.parameters():
param.requires_grad = True
net3.z2_h.requires_grad = False # z2_h constant
loss3.backward(retain_graph=True)
optimizer3.step()
for param in net3.parameters():
param.requires_grad = False
net3.z2_h.requires_grad = True # W and bias contant
loss3.backward()
optimizer3.step()
But the problem is that the loss3 does not decrease after the first few iterations and the network won’t learn finally (classification accuracy of ~10% on Cifar-10)! I’m not very familiar with Pytorch and I was wondering if you guys have any thoughts on this. Thanks!