Update xmisc
This commit is contained in:
		| @@ -17,6 +17,6 @@ kwargs: | ||||
|           module_path: torchvision.transforms | ||||
|           args: [] | ||||
|           kwargs: | ||||
|             mean: (0.491, 0.482, 0.447) | ||||
|             std: (0.247, 0.244, 0.262) | ||||
|             mean: [0.491, 0.482, 0.447] | ||||
|             std: [0.247, 0.244, 0.262] | ||||
|     kwargs: {} | ||||
|   | ||||
| @@ -25,6 +25,6 @@ kwargs: | ||||
|           module_path: torchvision.transforms | ||||
|           args: [] | ||||
|           kwargs: | ||||
|             mean: (0.491, 0.482, 0.447) | ||||
|             std: (0.247, 0.244, 0.262) | ||||
|             mean: [0.491, 0.482, 0.447] | ||||
|             std: [0.247, 0.244, 0.262] | ||||
|     kwargs: {} | ||||
|   | ||||
| @@ -58,6 +58,7 @@ def main(args): | ||||
|         pin_memory=True, | ||||
|         drop_last=False, | ||||
|     ) | ||||
|     iters_per_epoch = len(train_data) // args.batch_size | ||||
|  | ||||
|     logger.log("The training loader: {:}".format(train_loader)) | ||||
|     logger.log("The validation loader: {:}".format(valid_loader)) | ||||
| @@ -67,159 +68,44 @@ def main(args): | ||||
|         lr=args.lr, | ||||
|         weight_decay=args.weight_decay, | ||||
|     ) | ||||
|     loss = xmisc.nested_call_by_yaml(args.loss_config) | ||||
|     objective = xmisc.nested_call_by_yaml(args.loss_config) | ||||
|  | ||||
|     logger.log("The optimizer is:\n{:}".format(optimizer)) | ||||
|     logger.log("The loss is {:}".format(loss)) | ||||
|     logger.log("The objective is {:}".format(objective)) | ||||
|     logger.log("The iters_per_epoch={:}".format(iters_per_epoch)) | ||||
|  | ||||
|     model, loss = torch.nn.DataParallel(model).cuda(), loss.cuda() | ||||
|     model, objective = torch.nn.DataParallel(model).cuda(), objective.cuda() | ||||
|     scheduler = xmisc.LRMultiplier( | ||||
|         optimizer, xmisc.get_scheduler(args.scheduler, args.lr), args.steps | ||||
|     ) | ||||
|  | ||||
|     import pdb | ||||
|  | ||||
|     pdb.set_trace() | ||||
|  | ||||
|     train_func, valid_func = get_procedures(args.procedure) | ||||
|  | ||||
|     total_epoch = optim_config.epochs + optim_config.warmup | ||||
|     # Main Training and Evaluation Loop | ||||
|     start_time = time.time() | ||||
|     epoch_time = AverageMeter() | ||||
|     for epoch in range(start_epoch, total_epoch): | ||||
|         scheduler.update(epoch, 0.0) | ||||
|     start_time, iter_time = time.time(), xmisc.AverageMeter() | ||||
|     for xiter, data in enumerate(train_loader): | ||||
|         need_time = "Time Left: {:}".format( | ||||
|             convert_secs2time(epoch_time.avg * (total_epoch - epoch), True) | ||||
|         ) | ||||
|         epoch_str = "epoch={:03d}/{:03d}".format(epoch, total_epoch) | ||||
|         LRs = scheduler.get_lr() | ||||
|         find_best = False | ||||
|         # set-up drop-out ratio | ||||
|         if hasattr(base_model, "update_drop_path"): | ||||
|             base_model.update_drop_path( | ||||
|                 model_config.drop_path_prob * epoch / total_epoch | ||||
|             ) | ||||
|         logger.log( | ||||
|             "\n***{:s}*** start {:s} {:s}, LR=[{:.6f} ~ {:.6f}], scheduler={:}".format( | ||||
|                 time_string(), epoch_str, need_time, min(LRs), max(LRs), scheduler | ||||
|             xmisc.time_utils.convert_secs2time( | ||||
|                 iter_time.avg * (len(train_loader) - xiter), True | ||||
|             ) | ||||
|         ) | ||||
|         iter_str = "{:6d}/{:06d}".format(xiter, len(train_loader)) | ||||
|  | ||||
|         # train for one epoch | ||||
|         train_loss, train_acc1, train_acc5 = train_func( | ||||
|             train_loader, | ||||
|             network, | ||||
|             criterion, | ||||
|             scheduler, | ||||
|             optimizer, | ||||
|             optim_config, | ||||
|             epoch_str, | ||||
|             args.print_freq, | ||||
|             logger, | ||||
|         ) | ||||
|         # log the results | ||||
|         logger.log( | ||||
|             "***{:s}*** TRAIN [{:}] loss = {:.6f}, accuracy-1 = {:.2f}, accuracy-5 = {:.2f}".format( | ||||
|                 time_string(), epoch_str, train_loss, train_acc1, train_acc5 | ||||
|             ) | ||||
|         ) | ||||
|         inputs, targets = data | ||||
|         targets = targets.cuda(non_blocking=True) | ||||
|         model.train() | ||||
|  | ||||
|         # evaluate the performance | ||||
|         if (epoch % args.eval_frequency == 0) or (epoch + 1 == total_epoch): | ||||
|             logger.log("-" * 150) | ||||
|             valid_loss, valid_acc1, valid_acc5 = valid_func( | ||||
|                 valid_loader, | ||||
|                 network, | ||||
|                 criterion, | ||||
|                 optim_config, | ||||
|                 epoch_str, | ||||
|                 args.print_freq_eval, | ||||
|                 logger, | ||||
|             ) | ||||
|             valid_accuracies[epoch] = valid_acc1 | ||||
|             logger.log( | ||||
|                 "***{:s}*** VALID [{:}] loss = {:.6f}, accuracy@1 = {:.2f}, accuracy@5 = {:.2f} | Best-Valid-Acc@1={:.2f}, Error@1={:.2f}".format( | ||||
|                     time_string(), | ||||
|                     epoch_str, | ||||
|                     valid_loss, | ||||
|                     valid_acc1, | ||||
|                     valid_acc5, | ||||
|                     valid_accuracies["best"], | ||||
|                     100 - valid_accuracies["best"], | ||||
|                 ) | ||||
|             ) | ||||
|             if valid_acc1 > valid_accuracies["best"]: | ||||
|                 valid_accuracies["best"] = valid_acc1 | ||||
|                 find_best = True | ||||
|                 logger.log( | ||||
|                     "Currently, the best validation accuracy found at {:03d}-epoch :: acc@1={:.2f}, acc@5={:.2f}, error@1={:.2f}, error@5={:.2f}, save into {:}.".format( | ||||
|                         epoch, | ||||
|                         valid_acc1, | ||||
|                         valid_acc5, | ||||
|                         100 - valid_acc1, | ||||
|                         100 - valid_acc5, | ||||
|                         model_best_path, | ||||
|                     ) | ||||
|                 ) | ||||
|             num_bytes = ( | ||||
|                 torch.cuda.max_memory_cached(next(network.parameters()).device) * 1.0 | ||||
|             ) | ||||
|             logger.log( | ||||
|                 "[GPU-Memory-Usage on {:} is {:} bytes, {:.2f} KB, {:.2f} MB, {:.2f} GB.]".format( | ||||
|                     next(network.parameters()).device, | ||||
|                     int(num_bytes), | ||||
|                     num_bytes / 1e3, | ||||
|                     num_bytes / 1e6, | ||||
|                     num_bytes / 1e9, | ||||
|                 ) | ||||
|             ) | ||||
|             max_bytes[epoch] = num_bytes | ||||
|         if epoch % 10 == 0: | ||||
|             torch.cuda.empty_cache() | ||||
|         optimizer.zero_grad() | ||||
|         outputs = model(inputs) | ||||
|         loss = objective(outputs, targets) | ||||
|  | ||||
|         # save checkpoint | ||||
|         save_path = save_checkpoint( | ||||
|             { | ||||
|                 "epoch": epoch, | ||||
|                 "args": deepcopy(args), | ||||
|                 "max_bytes": deepcopy(max_bytes), | ||||
|                 "FLOP": flop, | ||||
|                 "PARAM": param, | ||||
|                 "valid_accuracies": deepcopy(valid_accuracies), | ||||
|                 "model-config": model_config._asdict(), | ||||
|                 "optim-config": optim_config._asdict(), | ||||
|                 "base-model": base_model.state_dict(), | ||||
|                 "scheduler": scheduler.state_dict(), | ||||
|                 "optimizer": optimizer.state_dict(), | ||||
|             }, | ||||
|             model_base_path, | ||||
|             logger, | ||||
|         ) | ||||
|         if find_best: | ||||
|             copy_checkpoint(model_base_path, model_best_path, logger) | ||||
|         last_info = save_checkpoint( | ||||
|             { | ||||
|                 "epoch": epoch, | ||||
|                 "args": deepcopy(args), | ||||
|                 "last_checkpoint": save_path, | ||||
|             }, | ||||
|             logger.path("info"), | ||||
|             logger, | ||||
|         ) | ||||
|         loss.backward() | ||||
|         optimizer.step() | ||||
|         scheduler.step() | ||||
|         if xiter % iters_per_epoch == 0: | ||||
|             logger.log("TRAIN [{:}] loss = {:.6f}".format(iter_str, loss.item())) | ||||
|  | ||||
|         # measure elapsed time | ||||
|         epoch_time.update(time.time() - start_time) | ||||
|         iter_time.update(time.time() - start_time) | ||||
|         start_time = time.time() | ||||
|  | ||||
|     logger.log("\n" + "-" * 200) | ||||
|     logger.log( | ||||
|         "Finish training/validation in {:} with Max-GPU-Memory of {:.2f} MB, and save final checkpoint into {:}".format( | ||||
|             convert_secs2time(epoch_time.sum, True), | ||||
|             max(v for k, v in max_bytes.items()) / 1e6, | ||||
|             logger.path("info"), | ||||
|         ) | ||||
|     ) | ||||
|     logger.log("-" * 200 + "\n") | ||||
|     logger.close() | ||||
|  | ||||
| @@ -249,7 +135,7 @@ if __name__ == "__main__": | ||||
|     parser.add_argument("--weight_decay", type=float, help="The weight decay") | ||||
|     parser.add_argument("--scheduler", type=str, help="The scheduler indicator.") | ||||
|     parser.add_argument("--steps", type=int, help="The total number of steps.") | ||||
|     parser.add_argument("--batch_size", type=int, default=2, help="The batch size.") | ||||
|     parser.add_argument("--batch_size", type=int, default=256, help="The batch size.") | ||||
|     parser.add_argument("--workers", type=int, default=4, help="The number of workers") | ||||
|     # Random Seed | ||||
|     parser.add_argument("--rand_seed", type=int, default=-1, help="manual seed") | ||||
|   | ||||
| @@ -28,4 +28,5 @@ python ./exps/basic/xmain.py --save_dir ${save_dir} --rand_seed ${rseed} \ | ||||
| 	--model_config ./configs/yaml.model/vit-cifar10.s0 \ | ||||
| 	--optim_config ./configs/yaml.opt/vit.cifar \ | ||||
| 	--loss_config ./configs/yaml.loss/cross-entropy \ | ||||
| 	--batch_size 256 \ | ||||
| 	--lr 0.003 --weight_decay 0.3 --scheduler warm-cos --steps 10000 | ||||
|   | ||||
| @@ -201,7 +201,6 @@ class SuperMLPv2(SuperModule): | ||||
|         self._hidden_multiplier = hidden_multiplier | ||||
|         self._out_features = out_features | ||||
|         self._drop_rate = drop | ||||
|         self._params = nn.ParameterDict({}) | ||||
|  | ||||
|         self._create_linear( | ||||
|             "fc1", self.in_features, int(self.in_features * self.hidden_multiplier) | ||||
| @@ -226,26 +225,22 @@ class SuperMLPv2(SuperModule): | ||||
|         return spaces.get_max(self._out_features) | ||||
|  | ||||
|     def _create_linear(self, name, inC, outC): | ||||
|         self._params["{:}_super_weight".format(name)] = torch.nn.Parameter( | ||||
|             torch.Tensor(outC, inC) | ||||
|         self.register_parameter( | ||||
|             "{:}_super_weight".format(name), torch.nn.Parameter(torch.Tensor(outC, inC)) | ||||
|         ) | ||||
|         self._params["{:}_super_bias".format(name)] = torch.nn.Parameter( | ||||
|             torch.Tensor(outC) | ||||
|         self.register_parameter( | ||||
|             "{:}_super_bias".format(name), torch.nn.Parameter(torch.Tensor(outC)) | ||||
|         ) | ||||
|  | ||||
|     def reset_parameters(self) -> None: | ||||
|         nn.init.kaiming_uniform_(self._params["fc1_super_weight"], a=math.sqrt(5)) | ||||
|         nn.init.kaiming_uniform_(self._params["fc2_super_weight"], a=math.sqrt(5)) | ||||
|         fan_in, _ = nn.init._calculate_fan_in_and_fan_out( | ||||
|             self._params["fc1_super_weight"] | ||||
|         ) | ||||
|         nn.init.kaiming_uniform_(self.fc1_super_weight, a=math.sqrt(5)) | ||||
|         nn.init.kaiming_uniform_(self.fc2_super_weight, a=math.sqrt(5)) | ||||
|         fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.fc1_super_weight) | ||||
|         bound = 1 / math.sqrt(fan_in) | ||||
|         nn.init.uniform_(self._params["fc1_super_bias"], -bound, bound) | ||||
|         fan_in, _ = nn.init._calculate_fan_in_and_fan_out( | ||||
|             self._params["fc2_super_weight"] | ||||
|         ) | ||||
|         nn.init.uniform_(self.fc1_super_bias, -bound, bound) | ||||
|         fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.fc2_super_weight) | ||||
|         bound = 1 / math.sqrt(fan_in) | ||||
|         nn.init.uniform_(self._params["fc2_super_bias"], -bound, bound) | ||||
|         nn.init.uniform_(self.fc2_super_bias, -bound, bound) | ||||
|  | ||||
|     @property | ||||
|     def abstract_search_space(self): | ||||
| @@ -282,8 +277,8 @@ class SuperMLPv2(SuperModule): | ||||
|         else: | ||||
|             hmul = spaces.get_determined_value(self._hidden_multiplier) | ||||
|         hidden_dim = int(expected_input_dim * hmul) | ||||
|         _fc1_weight = self._params["fc1_super_weight"][:hidden_dim, :expected_input_dim] | ||||
|         _fc1_bias = self._params["fc1_super_bias"][:hidden_dim] | ||||
|         _fc1_weight = self.fc1_super_weight[:hidden_dim, :expected_input_dim] | ||||
|         _fc1_bias = self.fc1_super_bias[:hidden_dim] | ||||
|         x = F.linear(input, _fc1_weight, _fc1_bias) | ||||
|         x = self.act(x) | ||||
|         x = self.drop(x) | ||||
| @@ -292,21 +287,17 @@ class SuperMLPv2(SuperModule): | ||||
|             out_dim = self.abstract_child["_out_features"].value | ||||
|         else: | ||||
|             out_dim = spaces.get_determined_value(self._out_features) | ||||
|         _fc2_weight = self._params["fc2_super_weight"][:out_dim, :hidden_dim] | ||||
|         _fc2_bias = self._params["fc2_super_bias"][:out_dim] | ||||
|         _fc2_weight = self.fc2_super_weight[:out_dim, :hidden_dim] | ||||
|         _fc2_bias = self.fc2_super_bias[:out_dim] | ||||
|         x = F.linear(x, _fc2_weight, _fc2_bias) | ||||
|         x = self.drop(x) | ||||
|         return x | ||||
|  | ||||
|     def forward_raw(self, input: torch.Tensor) -> torch.Tensor: | ||||
|         x = F.linear( | ||||
|             input, self._params["fc1_super_weight"], self._params["fc1_super_bias"] | ||||
|         ) | ||||
|         x = F.linear(input, self.fc1_super_weight, self.fc1_super_bias) | ||||
|         x = self.act(x) | ||||
|         x = self.drop(x) | ||||
|         x = F.linear( | ||||
|             x, self._params["fc2_super_weight"], self._params["fc2_super_bias"] | ||||
|         ) | ||||
|         x = F.linear(x, self.fc2_super_weight, self.fc2_super_bias) | ||||
|         x = self.drop(x) | ||||
|         return x | ||||
|  | ||||
|   | ||||
| @@ -1,319 +0,0 @@ | ||||
| ##################################################### | ||||
| # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.03 # | ||||
| ##################################################### | ||||
| import torch | ||||
| import torch.nn as nn | ||||
| import torch.nn.functional as F | ||||
|  | ||||
| import math | ||||
| from typing import Optional, Callable | ||||
|  | ||||
| from xautodl import spaces | ||||
| from .super_module import SuperModule | ||||
| from .super_module import IntSpaceType | ||||
| from .super_module import BoolSpaceType | ||||
|  | ||||
|  | ||||
| class SuperLinear(SuperModule): | ||||
|     """Applies a linear transformation to the incoming data: :math:`y = xA^T + b`""" | ||||
|  | ||||
|     def __init__( | ||||
|         self, | ||||
|         in_features: IntSpaceType, | ||||
|         out_features: IntSpaceType, | ||||
|         bias: BoolSpaceType = True, | ||||
|     ) -> None: | ||||
|         super(SuperLinear, self).__init__() | ||||
|  | ||||
|         # the raw input args | ||||
|         self._in_features = in_features | ||||
|         self._out_features = out_features | ||||
|         self._bias = bias | ||||
|         # weights to be optimized | ||||
|         self.register_parameter( | ||||
|             "_super_weight", | ||||
|             torch.nn.Parameter(torch.Tensor(self.out_features, self.in_features)), | ||||
|         ) | ||||
|         if self.bias: | ||||
|             self.register_parameter( | ||||
|                 "_super_bias", torch.nn.Parameter(torch.Tensor(self.out_features)) | ||||
|             ) | ||||
|         else: | ||||
|             self.register_parameter("_super_bias", None) | ||||
|         self.reset_parameters() | ||||
|  | ||||
|     @property | ||||
|     def in_features(self): | ||||
|         return spaces.get_max(self._in_features) | ||||
|  | ||||
|     @property | ||||
|     def out_features(self): | ||||
|         return spaces.get_max(self._out_features) | ||||
|  | ||||
|     @property | ||||
|     def bias(self): | ||||
|         return spaces.has_categorical(self._bias, True) | ||||
|  | ||||
|     @property | ||||
|     def abstract_search_space(self): | ||||
|         root_node = spaces.VirtualNode(id(self)) | ||||
|         if not spaces.is_determined(self._in_features): | ||||
|             root_node.append( | ||||
|                 "_in_features", self._in_features.abstract(reuse_last=True) | ||||
|             ) | ||||
|         if not spaces.is_determined(self._out_features): | ||||
|             root_node.append( | ||||
|                 "_out_features", self._out_features.abstract(reuse_last=True) | ||||
|             ) | ||||
|         if not spaces.is_determined(self._bias): | ||||
|             root_node.append("_bias", self._bias.abstract(reuse_last=True)) | ||||
|         return root_node | ||||
|  | ||||
|     def reset_parameters(self) -> None: | ||||
|         nn.init.kaiming_uniform_(self._super_weight, a=math.sqrt(5)) | ||||
|         if self.bias: | ||||
|             fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self._super_weight) | ||||
|             bound = 1 / math.sqrt(fan_in) | ||||
|             nn.init.uniform_(self._super_bias, -bound, bound) | ||||
|  | ||||
|     def forward_candidate(self, input: torch.Tensor) -> torch.Tensor: | ||||
|         # check inputs -> | ||||
|         if not spaces.is_determined(self._in_features): | ||||
|             expected_input_dim = self.abstract_child["_in_features"].value | ||||
|         else: | ||||
|             expected_input_dim = spaces.get_determined_value(self._in_features) | ||||
|         if input.size(-1) != expected_input_dim: | ||||
|             raise ValueError( | ||||
|                 "Expect the input dim of {:} instead of {:}".format( | ||||
|                     expected_input_dim, input.size(-1) | ||||
|                 ) | ||||
|             ) | ||||
|         # create the weight matrix | ||||
|         if not spaces.is_determined(self._out_features): | ||||
|             out_dim = self.abstract_child["_out_features"].value | ||||
|         else: | ||||
|             out_dim = spaces.get_determined_value(self._out_features) | ||||
|         candidate_weight = self._super_weight[:out_dim, :expected_input_dim] | ||||
|         # create the bias matrix | ||||
|         if not spaces.is_determined(self._bias): | ||||
|             if self.abstract_child["_bias"].value: | ||||
|                 candidate_bias = self._super_bias[:out_dim] | ||||
|             else: | ||||
|                 candidate_bias = None | ||||
|         else: | ||||
|             if spaces.get_determined_value(self._bias): | ||||
|                 candidate_bias = self._super_bias[:out_dim] | ||||
|             else: | ||||
|                 candidate_bias = None | ||||
|         return F.linear(input, candidate_weight, candidate_bias) | ||||
|  | ||||
|     def forward_raw(self, input: torch.Tensor) -> torch.Tensor: | ||||
|         return F.linear(input, self._super_weight, self._super_bias) | ||||
|  | ||||
|     def extra_repr(self) -> str: | ||||
|         return "in_features={:}, out_features={:}, bias={:}".format( | ||||
|             self._in_features, self._out_features, self._bias | ||||
|         ) | ||||
|  | ||||
|     def forward_with_container(self, input, container, prefix=[]): | ||||
|         super_weight_name = ".".join(prefix + ["_super_weight"]) | ||||
|         super_weight = container.query(super_weight_name) | ||||
|         super_bias_name = ".".join(prefix + ["_super_bias"]) | ||||
|         if container.has(super_bias_name): | ||||
|             super_bias = container.query(super_bias_name) | ||||
|         else: | ||||
|             super_bias = None | ||||
|         return F.linear(input, super_weight, super_bias) | ||||
|  | ||||
|  | ||||
| class SuperMLPv1(SuperModule): | ||||
|     """An MLP layer: FC -> Activation -> Drop -> FC -> Drop.""" | ||||
|  | ||||
|     def __init__( | ||||
|         self, | ||||
|         in_features: IntSpaceType, | ||||
|         hidden_features: IntSpaceType, | ||||
|         out_features: IntSpaceType, | ||||
|         act_layer: Callable[[], nn.Module] = nn.GELU, | ||||
|         drop: Optional[float] = None, | ||||
|     ): | ||||
|         super(SuperMLPv1, self).__init__() | ||||
|         self._in_features = in_features | ||||
|         self._hidden_features = hidden_features | ||||
|         self._out_features = out_features | ||||
|         self._drop_rate = drop | ||||
|         self.fc1 = SuperLinear(in_features, hidden_features) | ||||
|         self.act = act_layer() | ||||
|         self.fc2 = SuperLinear(hidden_features, out_features) | ||||
|         self.drop = nn.Dropout(drop or 0.0) | ||||
|  | ||||
|     @property | ||||
|     def abstract_search_space(self): | ||||
|         root_node = spaces.VirtualNode(id(self)) | ||||
|         space_fc1 = self.fc1.abstract_search_space | ||||
|         space_fc2 = self.fc2.abstract_search_space | ||||
|         if not spaces.is_determined(space_fc1): | ||||
|             root_node.append("fc1", space_fc1) | ||||
|         if not spaces.is_determined(space_fc2): | ||||
|             root_node.append("fc2", space_fc2) | ||||
|         return root_node | ||||
|  | ||||
|     def apply_candidate(self, abstract_child: spaces.VirtualNode): | ||||
|         super(SuperMLPv1, self).apply_candidate(abstract_child) | ||||
|         if "fc1" in abstract_child: | ||||
|             self.fc1.apply_candidate(abstract_child["fc1"]) | ||||
|         if "fc2" in abstract_child: | ||||
|             self.fc2.apply_candidate(abstract_child["fc2"]) | ||||
|  | ||||
|     def forward_candidate(self, input: torch.Tensor) -> torch.Tensor: | ||||
|         return self.forward_raw(input) | ||||
|  | ||||
|     def forward_raw(self, input: torch.Tensor) -> torch.Tensor: | ||||
|         x = self.fc1(input) | ||||
|         x = self.act(x) | ||||
|         x = self.drop(x) | ||||
|         x = self.fc2(x) | ||||
|         x = self.drop(x) | ||||
|         return x | ||||
|  | ||||
|     def extra_repr(self) -> str: | ||||
|         return "in_features={:}, hidden_features={:}, out_features={:}, drop={:}, fc1 -> act -> drop -> fc2 -> drop,".format( | ||||
|             self._in_features, | ||||
|             self._hidden_features, | ||||
|             self._out_features, | ||||
|             self._drop_rate, | ||||
|         ) | ||||
|  | ||||
|  | ||||
| class SuperMLPv2(SuperModule): | ||||
|     """An MLP layer: FC -> Activation -> Drop -> FC -> Drop.""" | ||||
|  | ||||
|     def __init__( | ||||
|         self, | ||||
|         in_features: IntSpaceType, | ||||
|         hidden_multiplier: IntSpaceType, | ||||
|         out_features: IntSpaceType, | ||||
|         act_layer: Callable[[], nn.Module] = nn.GELU, | ||||
|         drop: Optional[float] = None, | ||||
|     ): | ||||
|         super(SuperMLPv2, self).__init__() | ||||
|         self._in_features = in_features | ||||
|         self._hidden_multiplier = hidden_multiplier | ||||
|         self._out_features = out_features | ||||
|         self._drop_rate = drop | ||||
|         self._params = nn.ParameterDict({}) | ||||
|  | ||||
|         self._create_linear( | ||||
|             "fc1", self.in_features, int(self.in_features * self.hidden_multiplier) | ||||
|         ) | ||||
|         self._create_linear( | ||||
|             "fc2", int(self.in_features * self.hidden_multiplier), self.out_features | ||||
|         ) | ||||
|         self.act = act_layer() | ||||
|         self.drop = nn.Dropout(drop or 0.0) | ||||
|         self.reset_parameters() | ||||
|  | ||||
|     @property | ||||
|     def in_features(self): | ||||
|         return spaces.get_max(self._in_features) | ||||
|  | ||||
|     @property | ||||
|     def hidden_multiplier(self): | ||||
|         return spaces.get_max(self._hidden_multiplier) | ||||
|  | ||||
|     @property | ||||
|     def out_features(self): | ||||
|         return spaces.get_max(self._out_features) | ||||
|  | ||||
|     def _create_linear(self, name, inC, outC): | ||||
|         self._params["{:}_super_weight".format(name)] = torch.nn.Parameter( | ||||
|             torch.Tensor(outC, inC) | ||||
|         ) | ||||
|         self._params["{:}_super_bias".format(name)] = torch.nn.Parameter( | ||||
|             torch.Tensor(outC) | ||||
|         ) | ||||
|  | ||||
|     def reset_parameters(self) -> None: | ||||
|         nn.init.kaiming_uniform_(self._params["fc1_super_weight"], a=math.sqrt(5)) | ||||
|         nn.init.kaiming_uniform_(self._params["fc2_super_weight"], a=math.sqrt(5)) | ||||
|         fan_in, _ = nn.init._calculate_fan_in_and_fan_out( | ||||
|             self._params["fc1_super_weight"] | ||||
|         ) | ||||
|         bound = 1 / math.sqrt(fan_in) | ||||
|         nn.init.uniform_(self._params["fc1_super_bias"], -bound, bound) | ||||
|         fan_in, _ = nn.init._calculate_fan_in_and_fan_out( | ||||
|             self._params["fc2_super_weight"] | ||||
|         ) | ||||
|         bound = 1 / math.sqrt(fan_in) | ||||
|         nn.init.uniform_(self._params["fc2_super_bias"], -bound, bound) | ||||
|  | ||||
|     @property | ||||
|     def abstract_search_space(self): | ||||
|         root_node = spaces.VirtualNode(id(self)) | ||||
|         if not spaces.is_determined(self._in_features): | ||||
|             root_node.append( | ||||
|                 "_in_features", self._in_features.abstract(reuse_last=True) | ||||
|             ) | ||||
|         if not spaces.is_determined(self._hidden_multiplier): | ||||
|             root_node.append( | ||||
|                 "_hidden_multiplier", self._hidden_multiplier.abstract(reuse_last=True) | ||||
|             ) | ||||
|         if not spaces.is_determined(self._out_features): | ||||
|             root_node.append( | ||||
|                 "_out_features", self._out_features.abstract(reuse_last=True) | ||||
|             ) | ||||
|         return root_node | ||||
|  | ||||
|     def forward_candidate(self, input: torch.Tensor) -> torch.Tensor: | ||||
|         # check inputs -> | ||||
|         if not spaces.is_determined(self._in_features): | ||||
|             expected_input_dim = self.abstract_child["_in_features"].value | ||||
|         else: | ||||
|             expected_input_dim = spaces.get_determined_value(self._in_features) | ||||
|         if input.size(-1) != expected_input_dim: | ||||
|             raise ValueError( | ||||
|                 "Expect the input dim of {:} instead of {:}".format( | ||||
|                     expected_input_dim, input.size(-1) | ||||
|                 ) | ||||
|             ) | ||||
|         # create the weight and bias matrix for fc1 | ||||
|         if not spaces.is_determined(self._hidden_multiplier): | ||||
|             hmul = self.abstract_child["_hidden_multiplier"].value * expected_input_dim | ||||
|         else: | ||||
|             hmul = spaces.get_determined_value(self._hidden_multiplier) | ||||
|         hidden_dim = int(expected_input_dim * hmul) | ||||
|         _fc1_weight = self._params["fc1_super_weight"][:hidden_dim, :expected_input_dim] | ||||
|         _fc1_bias = self._params["fc1_super_bias"][:hidden_dim] | ||||
|         x = F.linear(input, _fc1_weight, _fc1_bias) | ||||
|         x = self.act(x) | ||||
|         x = self.drop(x) | ||||
|         # create the weight and bias matrix for fc2 | ||||
|         if not spaces.is_determined(self._out_features): | ||||
|             out_dim = self.abstract_child["_out_features"].value | ||||
|         else: | ||||
|             out_dim = spaces.get_determined_value(self._out_features) | ||||
|         _fc2_weight = self._params["fc2_super_weight"][:out_dim, :hidden_dim] | ||||
|         _fc2_bias = self._params["fc2_super_bias"][:out_dim] | ||||
|         x = F.linear(x, _fc2_weight, _fc2_bias) | ||||
|         x = self.drop(x) | ||||
|         return x | ||||
|  | ||||
|     def forward_raw(self, input: torch.Tensor) -> torch.Tensor: | ||||
|         x = F.linear( | ||||
|             input, self._params["fc1_super_weight"], self._params["fc1_super_bias"] | ||||
|         ) | ||||
|         x = self.act(x) | ||||
|         x = self.drop(x) | ||||
|         x = F.linear( | ||||
|             x, self._params["fc2_super_weight"], self._params["fc2_super_bias"] | ||||
|         ) | ||||
|         x = self.drop(x) | ||||
|         return x | ||||
|  | ||||
|     def extra_repr(self) -> str: | ||||
|         return "in_features={:}, hidden_multiplier={:}, out_features={:}, drop={:}, fc1 -> act -> drop -> fc2 -> drop,".format( | ||||
|             self._in_features, | ||||
|             self._hidden_multiplier, | ||||
|             self._out_features, | ||||
|             self._drop_rate, | ||||
|         ) | ||||
| @@ -1,6 +1,7 @@ | ||||
| ##################################################### | ||||
| # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.06 # | ||||
| ##################################################### | ||||
| """The module and yaml related functions.""" | ||||
| from .module_utils import call_by_dict | ||||
| from .module_utils import call_by_yaml | ||||
| from .module_utils import nested_call_by_dict | ||||
| @@ -11,10 +12,13 @@ from .torch_utils import count_parameters | ||||
|  | ||||
| from .logger_utils import Logger | ||||
|  | ||||
| # sampler | ||||
| """The data sampler related classes.""" | ||||
| from .sampler_utils import BatchSampler | ||||
|  | ||||
| # scheduler related | ||||
| """The meter related classes.""" | ||||
| from .meter_utils import AverageMeter | ||||
|  | ||||
| """The scheduler related classes.""" | ||||
| from .scheduler_utils import CosineParamScheduler, WarmupParamScheduler, LRMultiplier | ||||
|  | ||||
|  | ||||
|   | ||||
							
								
								
									
										22
									
								
								xautodl/xmisc/meter_utils.py
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										22
									
								
								xautodl/xmisc/meter_utils.py
									
									
									
									
									
										Normal file
									
								
							| @@ -0,0 +1,22 @@ | ||||
| class AverageMeter: | ||||
|     """Computes and stores the average and current value""" | ||||
|  | ||||
|     def __init__(self): | ||||
|         self.reset() | ||||
|  | ||||
|     def reset(self): | ||||
|         self.val = 0.0 | ||||
|         self.avg = 0.0 | ||||
|         self.sum = 0.0 | ||||
|         self.count = 0.0 | ||||
|  | ||||
|     def update(self, val, n=1): | ||||
|         self.val = val | ||||
|         self.sum += val * n | ||||
|         self.count += n | ||||
|         self.avg = self.sum / self.count | ||||
|  | ||||
|     def __repr__(self): | ||||
|         return "{name}(val={val}, avg={avg}, count={count})".format( | ||||
|             name=self.__class__.__name__, **self.__dict__ | ||||
|         ) | ||||
		Reference in New Issue
	
	Block a user