import torch import torch.nn as nn class AlexNet(nn.Module): q = 5 def __init__(self): super(AlexNet, self).__init__() self.features = nn.Sequential( nn.Conv2d(3, 8, kernel_size=3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(8, 1, kernel_size=3, padding=1), nn.ReLU(inplace=True), ) def forward(self, x): x = self.features(x) return x class My_loss(nn.Module): def __init__(self): super().__init__() def forward(self, x, y): return torch.mean(torch.pow((x - y), 2)) if __name__ == '__main__': criterion = My_loss() model = AlexNet() optimizer = torch.optim.SGD(params=model.parameters(), lr=0.0001, weight_decay=0.0001) for i in range(1000000): inputs = torch.randn((4, 3, 8, 8)) outputs = model(inputs) targets = torch.ones((4, 1, 8, 8)) loss = criterion(outputs, targets) loss.backward() # 计算梯度 optimizer.step() # 更新网络参数 optimizer.zero_grad() if i % 10 == 0: test = torch.randn((4, 3, 8, 8)) re = model(test) print((torch.mean(torch.abs(re)).item()))