Hi,
I tried adding torch.Variable to my custom (nn.module) and add it to optimizer group.
Variable name is ‘self.scale’
Code is like below.
class K32_BatchNet(nn.Module):
def __init__(self):
super(K32_BatchNet, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=5, padding=2),
nn.BatchNorm2d(32))
self.scale = Variable(torch.ones(1).cuda(), requires_grad=True)
self.layer2 = nn.Sequential(
nn.Conv2d(32, 64, kernel_size=5, padding=2),
nn.BatchNorm2d(64))
self.layer3 = nn.Sequential(
nn.Conv2d(64, 128, kernel_size=5, padding=2),
nn.BatchNorm2d(128))
self.conv3_drop = nn.Dropout2d()
self.avgPool2d = nn.AvgPool2d(3, stride=2)
self.fc1 = nn.Linear(1152,160)
self.fc2 = nn.Linear(160, 10)
def forward(self, x):
out = F.relu(self.layer1(x))
out = CustomFunction()(out, self.scale)
out = F.max_pool2d(out, 2)
Then, in main function,
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=5e-4)
optimizer.add_param_group({"params": model.scale})
In train function which is called in the main function,
data, target = Variable(data), Variable(target)
optimizer.zero_grad()
output = model(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
However, it shows below error.
Traceback (most recent call last):
File "k32_train.py", line 298, in <module>
main()
File "k32_train.py", line 93, in main
train(train_loader, model, optimizer, epoch)
File "k32_train.py", line 231, in train
optimizer.step()
File "/SSD/root/.local/lib/python2.7/site-packages/torch/optim/sgd.py", line 93, in step
d_p.add_(weight_decay, p.data)
RuntimeError: expand(torch.cuda.FloatTensor{[1]}, size=[]): the number of sizes provided (0) must be greater or equal to the number of dimensions in the tensor (1)
Please help me to get out of this problem.
Thanks!