Add files via upload

ShiYunyi · web-flow · commit 449c6b60a39e · 2021-08-28T22:19:58.000+08:00
diff --git a/README.md b/README.md
@@ -0,0 +1,36 @@
+# Meta-Weight-Net Code Optimization
+A new code framework that uses pytorch to implement meta-learning, and takes Meta-Weight-Net as an example.
+
+---
+
+By using a trick, meta-learning and meta-networks have become plug-and-play. We can now apply the meta learning
+algorithm directly to the existing pytorch model without rewriting it. 
+
+This code takes Meta-Weight-Net ([Meta-Weight-Net: Learning an Explicit Mapping For Sample Weighting](https://arxiv.org/abs/1902.07379))
+as an example to show how to use this trick. It rewrites an optimizer to assign non leaf node tensors to model parameters.
+See `meta.py` and line 90-120 of `main.py` for details.
+
+---
+##Environment
+- python 3.8
+- pytorch 1.9.0
+- torchvision 0.10.0
+
+`noisy_long_tail_CIFAR.py` can generate noisy and long-tailed CIFAR datasets by calling `torchvision.datasets`. Because 
+some class attributes have been changed, errors may occur in some earlier versions of torchvision. It can be solved by
+changing the corresponding attribute name.
+---
+##Running this example
+ResNet32 on CIFAR10-LT with imbalanced factor of 50:
+```
+python main.py --imbalanced_factor 50
+```
+ResNet32 on CIFAR10 with 40% uniform noise:
+```
+python main.py --meta_lr 1e-3 --meta_weight_decay 1e-4 --corruption_type uniform --corruption_ratio 0.4
+```
+---
+##Acknowledgements
+Thanks to the original code of Meta-Weight-Net (https://github.com/xjtushujun/meta-weight-net).
+
+Contact: Shi Yunyi (2404208668@qq.com)
diff --git a/main.py b/main.py
@@ -0,0 +1,163 @@
+import argparse
+import torch.optim
+# from torch.utils.tensorboard import SummaryWriter
+from meta import *
+from model import *
+from noisy_long_tail_CIFAR import *
+from utils import *
+
+
+parser = argparse.ArgumentParser(description='Meta_Weight_Net')
+parser.add_argument('--device', type=str, default='cuda')
+parser.add_argument('--seed', type=int, default=1)
+parser.add_argument('--meta_net_hidden_size', type=int, default=100)
+parser.add_argument('--meta_net_num_layers', type=int, default=1)
+
+parser.add_argument('--lr', type=float, default=.1)
+parser.add_argument('--momentum', type=float, default=.9)
+parser.add_argument('--dampening', type=float, default=0.)
+parser.add_argument('--nesterov', type=bool, default=False)
+parser.add_argument('--weight_decay', type=float, default=5e-4)
+parser.add_argument('--meta_lr', type=float, default=1e-5)
+parser.add_argument('--meta_weight_decay', type=float, default=0.)
+
+parser.add_argument('--dataset', type=str, default='cifar10')
+parser.add_argument('--num_meta', type=int, default=1000)
+parser.add_argument('--imbalanced_factor', type=int, default=None)
+parser.add_argument('--corruption_type', type=str, default=None)
+parser.add_argument('--corruption_ratio', type=float, default=0.)
+parser.add_argument('--batch_size', type=int, default=100)
+parser.add_argument('--max_epoch', type=int, default=120)
+
+parser.add_argument('--meta_interval', type=int, default=1)
+parser.add_argument('--paint_interval', type=int, default=20)
+
+args = parser.parse_args()
+print(args)
+
+
+def meta_weight_net():
+    set_cudnn(device=args.device)
+    set_seed(seed=args.seed)
+#     writer = SummaryWriter(log_dir='.\\mwn')
+
+    meta_net = MLP(hidden_size=args.meta_net_hidden_size, num_layers=args.meta_net_num_layers).to(device=args.device)
+    net = ResNet32(args.dataset == 'cifar10' and 10 or 100).to(device=args.device)
+
+    criterion = nn.CrossEntropyLoss().to(device=args.device)
+
+    optimizer = torch.optim.SGD(
+        net.parameters(),
+        lr=args.lr,
+        momentum=args.momentum,
+        dampening=args.dampening,
+        weight_decay=args.weight_decay,
+        nesterov=args.nesterov,
+    )
+    meta_optimizer = torch.optim.Adam(meta_net.parameters(), lr=args.meta_lr, weight_decay=args.meta_weight_decay)
+    lr = args.lr
+
+    train_dataloader, meta_dataloader, test_dataloader, imbalanced_num_list = build_dataloader(
+        seed=args.seed,
+        dataset=args.dataset,
+        num_meta_total=args.num_meta,
+        imbalanced_factor=args.imbalanced_factor,
+        corruption_type=args.corruption_type,
+        corruption_ratio=args.corruption_ratio,
+        batch_size=args.batch_size,
+    )
+
+    meta_dataloader_iter = iter(meta_dataloader)
+#     with torch.no_grad():
+#         for point in range(500):
+#             x = torch.tensor(point / 10).unsqueeze(0).to(args.device)
+#             fx = meta_net(x)
+#             writer.add_scalar('Initial Meta Net', fx, point)
+
+    for epoch in range(args.max_epoch):
+
+        if epoch >= 80 and epoch % 20 == 0:
+            lr = lr / 10
+        for group in optimizer.param_groups:
+            group['lr'] = lr
+
+        print('Training...')
+        for iteration, (inputs, labels) in enumerate(train_dataloader):
+            net.train()
+            inputs, labels = inputs.to(args.device), labels.to(args.device)
+
+            if (iteration + 1) % args.meta_interval == 0:
+                pseudo_net = ResNet32(args.dataset == 'cifar10' and 10 or 100).to(args.device)
+                pseudo_net.load_state_dict(net.state_dict())
+                pseudo_net.train()
+
+                pseudo_outputs = pseudo_net(inputs)
+                pseudo_loss_vector = functional.cross_entropy(pseudo_outputs, labels.long(), reduction='none')
+                pseudo_loss_vector_reshape = torch.reshape(pseudo_loss_vector, (-1, 1))
+                pseudo_weight = meta_net(pseudo_loss_vector_reshape.data)
+                pseudo_loss = torch.mean(pseudo_weight * pseudo_loss_vector_reshape)
+
+                pseudo_grads = torch.autograd.grad(pseudo_loss, pseudo_net.parameters(), create_graph=True)
+
+                pseudo_optimizer = MetaSGD(pseudo_net, pseudo_net.parameters(), lr=lr)
+                pseudo_optimizer.load_state_dict(optimizer.state_dict())
+                pseudo_optimizer.meta_step(pseudo_grads)
+
+                del pseudo_grads
+
+                try:
+                    meta_inputs, meta_labels = next(meta_dataloader_iter)
+                except StopIteration:
+                    meta_dataloader_iter = iter(meta_dataloader)
+                    meta_inputs, meta_labels = next(meta_dataloader_iter)
+
+                meta_inputs, meta_labels = meta_inputs.to(args.device), meta_labels.to(args.device)
+                meta_outputs = pseudo_net(meta_inputs)
+                meta_loss = criterion(meta_outputs, meta_labels.long())
+
+                meta_optimizer.zero_grad()
+                meta_loss.backward()
+                meta_optimizer.step()
+
+            outputs = net(inputs)
+            loss_vector = functional.cross_entropy(outputs, labels.long(), reduction='none')
+            loss_vector_reshape = torch.reshape(loss_vector, (-1, 1))
+
+            with torch.no_grad():
+                weight = meta_net(loss_vector_reshape)
+
+            loss = torch.mean(weight * loss_vector_reshape)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+        print('Computing Test Result...')
+        test_loss, test_accuracy = compute_loss_accuracy(
+            net=net,
+            data_loader=test_dataloader,
+            criterion=criterion,
+            device=args.device,
+        )
+#         writer.add_scalar('Loss', test_loss, epoch)
+#         writer.add_scalar('Accuracy', test_accuracy, epoch)
+
+        print('Epoch: {}, (Loss, Accuracy) Test: ({:.4f}, {:.2%}) LR: {}'.format(
+            epoch,
+            test_loss,
+            test_accuracy,
+            lr,
+        ))
+
+#         if (epoch + 1) % args.paint_interval == 0:
+#             with torch.no_grad():
+#                 for point in range(500):
+#                     x = torch.tensor(point / 10).unsqueeze(0).to(args.device)
+#                     fx = meta_net(x)
+#                     writer.add_scalar('Meta Net of Epoch {}'.format(epoch), fx, point)
+
+#     writer.close()
+
+
+if __name__ == '__main__':
+    meta_weight_net()
diff --git a/meta.py b/meta.py
@@ -0,0 +1,44 @@
+from torch.optim.sgd import SGD
+
+
+class MetaSGD(SGD):
+    def __init__(self, net, *args, **kwargs):
+        super(MetaSGD, self).__init__(*args, **kwargs)
+        self.net = net
+
+    def set_parameter(self, current_module, name, parameters):
+        if '.' in name:
+            name_split = name.split('.')
+            module_name = name_split[0]
+            rest_name = '.'.join(name_split[1:])
+            for children_name, children in current_module.named_children():
+                if module_name == children_name:
+                    self.set_parameter(children, rest_name, parameters)
+                    break
+        else:
+            current_module._parameters[name] = parameters
+
+    def meta_step(self, grads):
+        group = self.param_groups[0]
+        weight_decay = group['weight_decay']
+        momentum = group['momentum']
+        dampening = group['dampening']
+        nesterov = group['nesterov']
+        lr = group['lr']
+
+        for (name, parameter), grad in zip(self.net.named_parameters(), grads):
+            parameter.detach_()
+            if weight_decay != 0:
+                grad_wd = grad.add(parameter, alpha=weight_decay)
+            else:
+                grad_wd = grad
+            if momentum != 0 and 'momentum_buffer' in self.state[parameter]:
+                buffer = self.state[parameter]['momentum_buffer']
+                grad_b = buffer.mul(momentum).add(grad_wd, alpha=1-dampening)
+            else:
+                grad_b = grad_wd
+            if nesterov:
+                grad_n = grad_wd.add(grad_b, alpha=momentum)
+            else:
+                grad_n = grad_b
+            self.set_parameter(self.net, name, parameter.add(grad_n, alpha=-lr))
diff --git a/model.py b/model.py
@@ -0,0 +1,108 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as functional
+import torch.nn.init as init
+
+
+def _weights_init(m):
+    if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d):
+        init.kaiming_normal_(m.weight)
+
+
+class LambdaLayer(nn.Module):
+    def __init__(self, lambd):
+        super(LambdaLayer, self).__init__()
+        self.lambd = lambd
+
+    def forward(self, x):
+        return self.lambd(x)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1, option='A'):
+        super(BasicBlock, self).__init__()
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != planes:
+            if option == 'A':
+                """
+                For CIFAR10 ResNet paper uses option A.
+                """
+                self.shortcut = LambdaLayer(lambda x:
+                                            functional.pad(x[:, :, ::2, ::2], (0, 0, 0, 0, planes//4, planes//4), "constant", 0))
+            elif option == 'B':
+                self.shortcut = nn.Sequential(
+                    nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
+                    nn.BatchNorm2d(self.expansion * planes)
+                )
+
+    def forward(self, x):
+        out = functional.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = functional.relu(out)
+        return out
+
+
+class ResNet32(nn.Module):
+    def __init__(self, num_classes=10, block=BasicBlock, num_blocks=[5, 5, 5]):
+        super(ResNet32, self).__init__()
+        self.in_planes = 16
+
+        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(16)
+        self.layer1 = self._make_layer(block, 16, num_blocks[0], stride=1)
+        self.layer2 = self._make_layer(block, 32, num_blocks[1], stride=2)
+        self.layer3 = self._make_layer(block, 64, num_blocks[2], stride=2)
+        self.linear = nn.Linear(64, num_classes)
+
+        self.apply(_weights_init)
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        strides = [stride] + [1]*(num_blocks-1)
+        layers = []
+        for stride in strides:
+            layers.append(block(self.in_planes, planes, stride))
+            self.in_planes = planes * block.expansion
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        out = functional.relu(self.bn1(self.conv1(x)))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = functional.avg_pool2d(out, out.size()[3])
+        out = out.view(out.size(0), -1)
+        out = self.linear(out)
+        return out
+
+
+class HiddenLayer(nn.Module):
+    def __init__(self, input_size, output_size):
+        super(HiddenLayer, self).__init__()
+        self.fc = nn.Linear(input_size, output_size)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        return self.relu(self.fc(x))
+
+
+class MLP(nn.Module):
+    def __init__(self, hidden_size=100, num_layers=1):
+        super(MLP, self).__init__()
+        self.first_hidden_layer = HiddenLayer(1, hidden_size)
+        self.rest_hidden_layers = nn.Sequential(*[HiddenLayer(hidden_size, hidden_size) for _ in range(num_layers - 1)])
+        self.output_layer = nn.Linear(hidden_size, 1)
+
+    def forward(self, x):
+        x = self.first_hidden_layer(x)
+        x = self.rest_hidden_layers(x)
+        x = self.output_layer(x)
+        return torch.sigmoid(x)
diff --git a/noisy_long_tail_CIFAR.py b/noisy_long_tail_CIFAR.py
diff --git a/utils.py b/utils.py
