Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- class Net(nn.Module):
- def __init__(self):
- super(Net, self).__init__()
- self.act_max = nn.Parameter(torch.Tensor([0]), requires_grad=True)
- self.conv1 = nn.Conv2d(3, 32, kernel_size=5)
- self.conv2 = nn.Conv2d(32, 64, kernel_size=5)
- self.pool = nn.MaxPool2d(2, 2)
- self.relu = nn.ReLU()
- self.linear = nn.Linear(64 * 5 * 5, 10)
- def forward(self, input):
- conv1 = self.conv1(input)
- pool1 = self.pool(conv1)
- relu1 = self.relu(pool1)
- relu1[relu1 > self.act_max] = self.act_max
- conv2 = self.conv2(relu1)
- pool2 = self.pool(conv2)
- relu2 = self.relu(pool2)
- relu2 = relu2.view(relu2.size(0), -1)
- linear = self.linear(relu2)
- return linear
- model = Net()
- torch.nn.init.kaiming_normal_(model.parameters)
- nn.init.constant(model.act_max, 1.0)
- model = model.cuda()
- optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
- for epoch in range(100):
- for i in range(1000):
- output = model(input)
- loss = nn.CrossEntropyLoss()(output, label)
- optimizer.zero_grad()
- loss.backward()
- optimizer.step()
- model.act_max.data = model.act_max.data - 0.001 * model.act_max.grad.data
Add Comment
Please, Sign In to add comment