[ { "path": "Backbone/CNN1D.py", "content": "import torch.nn as nn\n\nclass CNN1D(nn.Module):\n def __init__(self, num_out = 10):\n super(CNN1D, self).__init__()\n self.layer1 = nn.Sequential(\n nn.Conv1d(1,32,kernel_size=3,padding=1),\n nn.BatchNorm1d(32),\n nn.ReLU(),\n nn.MaxPool1d(kernel_size=2, padding=0)\n )\n self.layer2 = nn.Sequential(\n nn.Conv1d(32,64,kernel_size=3,padding=1),\n nn.BatchNorm1d(64),\n nn.ReLU(),\n nn.MaxPool1d(kernel_size=2, padding=0)\n )\n self.layer3 = nn.Sequential(\n nn.Conv1d(64,64,kernel_size=3,padding=1),\n nn.BatchNorm1d(64),\n nn.ReLU(),\n nn.MaxPool1d(kernel_size=2, padding=0)\n )\n self.avgpool = nn.AdaptiveAvgPool1d(1) # output (64,1)\n self.fc = nn.Sequential(nn.Linear(64,num_out, nn.Dropout(0.5)))\n\n def forward(self, x):\n x = self.layer1(x)\n x = self.layer2(x)\n x = self.layer3(x)\n x = self.avgpool(x)\n x = x.view(-1,64)\n\n return x" }, { "path": "Backbone/MLPNet.py", "content": "import torch.nn as nn\n\nclass MLPNet(nn.Module):\n def __init__(self, num_in = 1024, num_out = 10):\n super(MLPNet, self).__init__()\n self.fc1 = nn.Sequential(\n nn.Linear(num_in,512),\n nn.BatchNorm1d(512),\n nn.ReLU()\n )\n self.fc2 = nn.Sequential(\n nn.Linear(512, 256),\n nn.BatchNorm1d(256),\n nn.ReLU()\n )\n self.fc3 = nn.Sequential(\n nn.Linear(256, 128),\n nn.BatchNorm1d(128),\n nn.ReLU()\n )\n self.fc4 = nn.Sequential(\n nn.Linear(128, 64),\n nn.BatchNorm1d(64),\n nn.ReLU()\n )\n self.fc5 = nn.Linear(64,num_out)\n\n def forward(self, x):\n x = self.fc1(x)\n x = self.fc2(x)\n x = self.fc3(x)\n x = self.fc4(x)\n\n return x" }, { "path": "Backbone/ResNet1D.py", "content": "# one-dimentional ResNet source code reference: https://github.com/ZhaoZhibin/UDTL/blob/master/models/resnet18_1d.py\n\nimport torch.nn as nn\nimport torch.utils.model_zoo as model_zoo\n\nmodel_urls = {\n 'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',\n 'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',\n 'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',\n 'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',\n 'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',\n}\n\ndef conv3x1(in_planes, out_planes, stride=1):\n \"\"\"3x3 convolution with padding\"\"\"\n return nn.Conv1d(in_planes, out_planes, kernel_size=3, stride=stride,\n padding=1, bias=False)\n\ndef conv1x1(in_planes, out_planes, stride=1):\n \"\"\"1x1 convolution\"\"\"\n return nn.Conv1d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)\n\n\nclass BasicBlock(nn.Module):\n expansion = 1\n\n def __init__(self, inplanes, planes, stride=1, downsample=None):\n super(BasicBlock, self).__init__()\n self.conv1 = conv3x1(inplanes, planes, stride)\n self.bn1 = nn.BatchNorm1d(planes)\n self.relu = nn.ReLU(inplace=True)\n self.conv2 = conv3x1(planes, planes)\n self.bn2 = nn.BatchNorm1d(planes)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n identity = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n\n if self.downsample is not None:\n identity = self.downsample(x)\n\n out += identity\n out = self.relu(out)\n\n return out\n\n\nclass Bottleneck(nn.Module):\n # For ResNet50, ResNet101, ResNet152\n expansion = 4\n\n def __init__(self, inplanes, planes, stride=1, downsample=None):\n super(Bottleneck, self).__init__()\n self.conv1 = conv1x1(inplanes, planes)\n self.bn1 = nn.BatchNorm1d(planes)\n self.conv2 = conv3x1(planes, planes, stride)\n self.bn2 = nn.BatchNorm1d(planes)\n self.conv3 = conv1x1(planes, planes * self.expansion)\n self.bn3 = nn.BatchNorm1d(planes * self.expansion)\n self.relu = nn.ReLU(inplace=True)\n self.downsample = downsample\n self.stride = stride\n\n def forward(self, x):\n identity = x\n\n out = self.conv1(x)\n out = self.bn1(out)\n out = self.relu(out)\n\n out = self.conv2(out)\n out = self.bn2(out)\n out = self.relu(out)\n\n out = self.conv3(out)\n out = self.bn3(out)\n\n if self.downsample is not None:\n identity = self.downsample(x)\n\n out += identity\n out = self.relu(out)\n\n return out\n\n\nclass ResNet(nn.Module):\n\n def __init__(self, block, layers, in_channel=1, out_channel=10, zero_init_residual=False):\n super(ResNet, self).__init__()\n self.inplanes = 64\n self.conv1 = nn.Conv1d(in_channel, 64, kernel_size=7, stride=2, padding=3,\n bias=False)\n self.bn1 = nn.BatchNorm1d(64)\n self.relu = nn.ReLU(inplace=True)\n self.maxpool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)\n self.layer1 = self._make_layer(block, 64, layers[0])\n self.layer2 = self._make_layer(block, 128, layers[1], stride=2)\n self.layer3 = self._make_layer(block, 256, layers[2], stride=2)\n self.layer4 = self._make_layer(block, 512, layers[3], stride=2)\n self.avgpool = nn.AdaptiveAvgPool1d(1) # output (512, 1)\n\n\n for m in self.modules():\n if isinstance(m, nn.Conv1d):\n nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n elif isinstance(m, nn.BatchNorm1d):\n nn.init.constant_(m.weight, 1)\n nn.init.constant_(m.bias, 0)\n\n # Zero-initialize the last BN in each residual branch,\n # so that the residual branch starts with zeros, and each residual block behaves like an identity.\n # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677\n if zero_init_residual:\n for m in self.modules():\n if isinstance(m, Bottleneck):\n nn.init.constant_(m.bn3.weight, 0)\n elif isinstance(m, BasicBlock):\n nn.init.constant_(m.bn2.weight, 0)\n\n def _make_layer(self, block, planes, blocks, stride=1):\n downsample = None\n if stride != 1 or self.inplanes != planes * block.expansion:\n downsample = nn.Sequential(\n conv1x1(self.inplanes, planes * block.expansion, stride),\n nn.BatchNorm1d(planes * block.expansion),\n )\n\n layers = []\n layers.append(block(self.inplanes, planes, stride, downsample))\n self.inplanes = planes * block.expansion\n for _ in range(1, blocks):\n layers.append(block(self.inplanes, planes))\n\n return nn.Sequential(*layers)\n\n def forward(self, x):\n x = self.conv1(x)\n x = self.bn1(x)\n x = self.relu(x)\n x = self.maxpool(x)\n\n x = self.layer1(x)\n x = self.layer2(x)\n x = self.layer3(x)\n x = self.layer4(x)\n\n x = self.avgpool(x)\n x = x.view(x.size(0), -1)\n\n return x\n\n\ndef resnet18(pretrained=False, **kwargs):\n \"\"\"Constructs a ResNet-18 model.\n\n Args:\n pretrained (bool): If True, returns a model pre-trained on ImageNet\n \"\"\"\n model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)\n if pretrained:\n model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))\n return model\n\n\n# convnet without the last layer\nclass resnet18_features(nn.Module):\n def __init__(self, pretrained=False):\n super(resnet18_features, self).__init__()\n self.model_resnet18 = resnet18(pretrained)\n self.__in_features = 512\n\n def forward(self, x):\n x = self.model_resnet18(x)\n return x\n\n def output_num(self):\n return self.__in_features" }, { "path": "DANN.py", "content": "import argparse\nimport os\nimport numpy as np\nfrom Utils.logger import setlogger\n\nfrom turtle import forward\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torchvision.models as models\n\nfrom Backbone import ResNet1D, MLPNet, CNN1D\nfrom loss import MKMMD, MMDLinear, CORAL\nfrom PreparData.CWRU import CWRUloader\nimport Utils.utils as utils\n\nfrom tqdm import *\nimport warnings\nimport logging\n\n# ===== Define argments =====\ndef parse_args():\n parser = argparse.ArgumentParser(description='Implementation of Domain Adversarial Neural Networks')\n\n # task setting\n parser.add_argument(\"--log_file\", type=str, default=\"./logs/DANN.log\", help=\"log file path\")\n\n # dataset information\n parser.add_argument(\"--datadir\", type=str, default=\"./datasets\", help=\"data directory\")\n parser.add_argument(\"--source_dataname\", type=str, default=\"CWRU\", choices=[\"CWRU\", \"PU\"], help=\"choice a dataset\")\n parser.add_argument(\"--target_dataname\", type=str, default=\"CWRU\", choices=[\"CWRU\", \"PU\"], help=\"choice a dataset\")\n parser.add_argument(\"--s_load\", type=int, default=3, help=\"source domain working condition\")\n parser.add_argument(\"--t_load\", type=int, default=2, help=\"target domain working condition\")\n parser.add_argument(\"--s_label_set\", type=list, default=[0,1,2,3,4,5,6,7,8,9], help=\"source domain label set\")\n parser.add_argument(\"--t_label_set\", type=list, default=[0,1,2,3,4,5,6,7,8,9], help=\"target domain label set\")\n parser.add_argument(\"--val_rat\", type=float, default=0.3, help=\"training-validation rate\")\n parser.add_argument(\"--test_rat\", type=float, default=0.5, help=\"validation-test rate\")\n parser.add_argument(\"--seed\", type=int, default=\"29\")\n\n # pre-processing\n parser.add_argument(\"--fft\", type=bool, default=False, help=\"FFT preprocessing\")\n parser.add_argument(\"--window\", type=int, default=128, help=\"time window, if not augment data, window=1024\")\n parser.add_argument(\"--normalization\", type=str, default=\"0-1\", choices=[\"None\", \"0-1\", \"mean-std\"], help=\"normalization option\")\n parser.add_argument(\"--savemodel\", type=bool, default=False, help=\"whether save pre-trained model in the classification task\")\n parser.add_argument(\"--pretrained\", type=bool, default=False, help=\"whether use pre-trained model in transfer learning tasks\")\n\n # backbone\n parser.add_argument(\"--backbone\", type=str, default=\"ResNet1D\", choices=[\"ResNet1D\", \"ResNet2D\", \"MLPNet\", \"CNN1D\"])\n # if backbone in (\"ResNet1D\", \"CNN1D\"), data shape: (batch size, 1, 1024)\n # elif backbone == \"ResNet2D\", data shape: (batch size, 3, 32, 32)\n # elif backbone == \"MLPNet\", data shape: (batch size, 1024)\n\n\n # optimization & training\n parser.add_argument(\"--num_workers\", type=int, default=0, help=\"the number of dataloader workers\")\n parser.add_argument(\"--batch_size\", type=int, default=256)\n parser.add_argument(\"--max_epoch\", type=int, default=100)\n parser.add_argument(\"--lr\", type=float, default=1e-3, help=\"learning rate\")\n parser.add_argument('--lr_scheduler', type=str, default='stepLR', choices=['step', 'exp', 'stepLR', 'fix'], help='the learning rate schedule')\n parser.add_argument('--gamma', type=float, default=0.8, help='learning rate scheduler parameter for step and exp')\n parser.add_argument('--steps', type=str, default='30, 120', help='the learning rate decay for step and stepLR')\n parser.add_argument(\"--optimizer\", type=str, default=\"adam\", choices=[\"adam\", \"sgd\"])\n\n\n args = parser.parse_args()\n return args\n\n# ===== Build Model =====\nclass FeatureNet(nn.Module):\n def __init__(self, args):\n super(FeatureNet, self).__init__()\n if args.backbone == \"ResNet1D\":\n self.feature_net = ResNet1D.resnet18()\n elif args.backbone == \"ResNet2D\":\n self.model_ft = models.resnet18(pretrained=True)\n self.bottleneck = nn.Sequential(nn.Linear(self.model_ft.fc.out_features, 512), nn.ReLU(), nn.Dropout(0.5))\n self.feature_net = nn.Sequential(self.model_ft, self.bottleneck)\n elif args.backbone == \"MLPNet\":\n if args.fft:\n self.feature_net = MLPNet.MLPNet(num_in=512)\n else:\n self.feature_net = MLPNet.MLPNet()\n elif args.backbone == \"CNN1D\":\n self.feature_net = CNN1D.CNN1D()\n else:\n raise Exception(\"model not implement\")\n\n def forward(self, x):\n logits = self.feature_net(x)\n\n return logits\n\nclass Classifier(nn.Module):\n def __init__(self, args, num_out=10):\n super(Classifier, self).__init__()\n if args.backbone in (\"ResNet1D\", \"ResNet2D\"):\n self.classifier = nn.Sequential(nn.Linear(512,num_out, nn.Dropout(0.5)))\n if args.backbone in (\"MLPNet\", \"CNN1D\"):\n self.classifier = nn.Sequential(nn.Linear(64,num_out, nn.Dropout(0.5)))\n\n def forward(self, logits):\n outputs = self.classifier(logits)\n\n return outputs\n\n# Define the discriminator\ndef calc_coeff(iter_num, high=1.0, low=0.0, alpha=10.0, max_iter=10000.0):\n return np.float(2.0 * (high - low) / (1.0 + np.exp(-alpha * iter_num / max_iter)) - (high - low) + low)\n\n## The hook will be called every time a gradient with respect to the Tensor is computed.\n## https://pytorch.org/docs/stable/generated/torch.Tensor.register_hook.html?highlight=register_hook#torch.Tensor.register_hook\ndef grl_hook(coeff):\n def fun1(grad):\n return -coeff * grad.clone()\n return fun1\n\nclass Discriminator(nn.Module):\n def __init__(self, args, num_out = 1, max_iter=10000.0, trade_off_adversarial='Cons', lam_adversarial=1.0):\n super(Discriminator, self).__init__()\n if args.backbone in (\"ResNet1D\", \"ResNet2D\"):\n self.domain_classifier = nn.Sequential(\n nn.Linear(512,128, nn.Dropout(0.5)),\n nn.BatchNorm1d(128),\n nn.ReLU(),\n nn.Linear(128, num_out)\n )\n elif args.backbone in (\"MLPNet\", \"CNN1D\"):\n self.domain_classifier = nn.Sequential(\n nn.Linear(64,32, nn.Dropout(0.5)),\n nn.BatchNorm1d(32),\n nn.ReLU(),\n nn.Linear(32, num_out)\n )\n self.sigmoid = nn.Sigmoid()\n\n # parameters\n self.iter_num = 0\n self.alpha = 10\n self.low = 0.0\n self.high = 1.0\n self.max_iter = max_iter\n self.trade_off_adversarial = trade_off_adversarial\n self.lam_adversarial = lam_adversarial\n \n def forward(self, x):\n if self.training:\n self.iter_num += 1\n if self.trade_off_adversarial == \"Cons\":\n coeff = self.lam_adversarial\n elif self.trade_off_adversarial == \"Step\":\n coeff = calc_coeff(self.iter_num, self.high, self.low,\\\n self.alpha, self.max_iter)\n else:\n raise Exception(\"loss not implement\")\n x = x * 1.0\n x.register_hook(grl_hook(coeff))\n x = self.domain_classifier(x)\n x = self.sigmoid(x)\n return x\n\n# ===== Load Data =====\ndef loaddata(args):\n if args.source_dataname == \"CWRU\":\n source_data, source_label = CWRUloader(args, args.s_load, args.s_label_set)\n\n source_data, source_label = np.concatenate(source_data, axis=0), np.concatenate(source_label, axis=0)\n \n if args.target_dataname == \"CWRU\":\n target_data, target_label = CWRUloader(args, args.t_load, args.t_label_set)\n\n target_data, target_label = np.concatenate(target_data, axis=0), np.concatenate(target_label, axis=0)\n\n source_loader, _, _ = utils.DataSplite(args, source_data, source_label)\n target_trainloader, target_valloader, target_testloader = utils.DataSplite(args, target_data, target_label)\n \n return source_loader, target_trainloader, target_valloader, target_testloader\n\n# ===== Test the Model =====\ndef tester(featurenet, classifier, dataloader):\n featurenet.eval()\n classifier.eval()\n device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n correct_num, total_num = 0, 0\n for i, (x_batch, y_batch) in enumerate(dataloader):\n x_batch, y_batch = x_batch.to(device), y_batch.to(device)\n\n # compute model cotput and loss\n logtis_batch = featurenet(x_batch)\n output_batch = classifier(logtis_batch)\n\n pre = torch.max(output_batch.cpu(), 1)[1].numpy()\n y = y_batch.cpu().numpy()\n correct_num += (pre == y).sum()\n total_num += len(y)\n accuracy = (correct_num / total_num) * 100.0\n return accuracy\n\n# ===== Train the Model =====\ndef trainer(args):\n # Consider the gpu or cpu condition\n if torch.cuda.is_available():\n device = torch.device(\"cuda\")\n device_count = torch.cuda.device_count()\n logging.info('using {} gpus'.format(device_count))\n assert args.batch_size % device_count == 0, \"batch size should be divided by device count\"\n else:\n warnings.warn(\"gpu is not available\")\n device = torch.device(\"cpu\")\n device_count = 1\n logging.info('using {} cpu'.format(device_count))\n \n # load the dataset\n source_trainloader, target_trainloader, target_valloader, target_testloader = loaddata(args)\n\n # load the model\n featurenet = FeatureNet(args)\n classifier = Classifier(args, num_out=len(args.t_label_set))\n discriminator = Discriminator(args)\n\n # load the checkpoint\n if args.pretrained:\n if args.backbone != \"ResNet2D\": # pretrained ResNet2D model is downloaded from torchvision module\n if not args.fft:\n path = \"./checkpoints/{}_checkpoint.tar\".format(args.backbone)\n else:\n path = \"./checkpoints/{}FFT_checkpoint.tar\".format(args.backbone)\n featurenet.load_state_dict(torch.load(path))\n\n parameter_list = [{\"params\": featurenet.parameters(), \"lr\": 0.5*args.lr},\n {\"params\": classifier.parameters(), \"lr\": args.lr},\n {\"params\": discriminator.parameters(), \"lr\": args.lr}]\n\n # Define optimizer and learning rate decay\n optimizer, lr_scheduler = utils.optimizer(args, parameter_list)\n\n ## define loss function\n loss_cls = nn.CrossEntropyLoss()\n loss_adver = nn.BCELoss()\n\n featurenet.to(device)\n classifier.to(device)\n discriminator.to(device)\n\n # train\n best_acc = 0.0\n meters = {\"acc_source_train\":[], \"acc_target_train\": [], \"acc_target_val\": []}\n\n for epoch in range(args.max_epoch):\n featurenet.train()\n classifier.train()\n with tqdm(total=len(target_trainloader), leave=False) as pbar:\n for i, ((x_s_batch, y_s_batch), (x_t_batch, y_t_batch)) in enumerate(zip(source_trainloader,target_trainloader)):\n\n if len(y_s_batch) != len(y_t_batch):\n break\n \n batch_num = x_s_batch.size(0)\n \n domain_label_source = torch.ones(batch_num).float()\n domain_label_target = torch.zeros(batch_num).float()\n\n inputs = torch.cat((x_s_batch, x_t_batch), dim=0)\n domain_label = torch.cat((domain_label_source, domain_label_target), dim=0)\n\n # move to GPU if available\n inputs = inputs.to(device)\n s_labels = y_s_batch.to(device)\n t_labels = y_t_batch.to(device)\n domain_label = domain_label.to(device)\n\n # compute model cotput and loss\n logits = featurenet(inputs)\n outputs = classifier(logits)\n domain_outputs = discriminator(logits)\n\n\n classification_loss = loss_cls(outputs.narrow(0, 0, batch_num), s_labels.long())\n adversarial_loss = loss_adver(domain_outputs.squeeze(), domain_label)\n loss = classification_loss + adversarial_loss\n\n\n # clear previous gradients, compute gradients\n optimizer.zero_grad()\n loss.backward()\n\n # performs updates using calculated gradients\n optimizer.step()\n\n # evaluate\n # training accuracy\n acc_source_train = utils.accuracy(outputs.narrow(0, 0, batch_num), s_labels)\n acc_target_train = utils.accuracy(outputs.narrow(0, batch_num, batch_num), t_labels) \n\n pbar.update()\n \n # update lr\n if lr_scheduler is not None:\n lr_scheduler.step()\n\n val_acc = tester(featurenet, classifier, target_valloader)\n if val_acc > best_acc:\n best_acc = val_acc\n if args.savemodel:\n utils.save_model(featurenet, args)\n \n logging.info(\"Epoch: {:>3}/{}, loss_cls: {:.4f}, loss: {:.4f}, source_train_acc: {:>6.2f}%, target_train_acc: {:>6.2f}%, target_val_acc: {:>6.2f}%\".format(\\\n epoch+1, args.max_epoch, classification_loss, loss, acc_source_train, acc_target_train, val_acc))\n meters[\"acc_source_train\"].append(acc_source_train)\n meters[\"acc_target_train\"].append(acc_target_train)\n meters[\"acc_target_val\"].append(val_acc)\n\n logging.info(\"Best accuracy: {:.4f}\".format(best_acc))\n utils.save_log(meters, \"./logs/DDC_{}_{}_meters.pkl\".format(args.backbone, args.max_epoch))\n\n logging.info(\"=\"*15+\"Done!\"+\"=\"*15)\n\nif __name__ == \"__main__\":\n\n args = parse_args()\n\n # set the logger\n if not os.path.exists(\"./logs\"):\n os.makedirs(\"./logs\")\n setlogger(args.log_file)\n\n # save the args\n for k, v in args.__dict__.items():\n logging.info(\"{}: {}\".format(k, v))\n \n trainer(args)" }, { "path": "DDC.py", "content": "import argparse\nimport os\nimport numpy as np\nfrom Utils.logger import setlogger\n\nfrom turtle import forward\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torchvision.models as models\n\nfrom Backbone import ResNet1D, MLPNet, CNN1D\nfrom loss import MKMMD, MMDLinear, CORAL\nfrom PreparData.CWRU import CWRUloader\nimport Utils.utils as utils\n\nfrom tqdm import *\nimport warnings\nimport logging\n\n# ===== Define argments =====\ndef parse_args():\n parser = argparse.ArgumentParser(description='Implementation of Deep Domain Confusion networks')\n\n # task setting\n parser.add_argument(\"--log_file\", type=str, default=\"./logs/DDC.log\", help=\"log file path\")\n\n # dataset information\n parser.add_argument(\"--datadir\", type=str, default=\"./datasets\", help=\"data directory\")\n parser.add_argument(\"--source_dataname\", type=str, default=\"CWRU\", choices=[\"CWRU\"], help=\"choice a dataset\")\n parser.add_argument(\"--target_dataname\", type=str, default=\"CWRU\", choices=[\"CWRU\"], help=\"choice a dataset\")\n parser.add_argument(\"--s_load\", type=int, default=3, help=\"source domain working condition\")\n parser.add_argument(\"--t_load\", type=int, default=2, help=\"target domain working condition\")\n parser.add_argument(\"--s_label_set\", type=list, default=[0,1,2,3,4,5,6,7,8,9], help=\"source domain label set\")\n parser.add_argument(\"--t_label_set\", type=list, default=[0,1,2,3,4,5,6,7,8,9], help=\"target domain label set\")\n parser.add_argument(\"--val_rat\", type=float, default=0.3, help=\"training-validation rate\")\n parser.add_argument(\"--test_rat\", type=float, default=0.5, help=\"validation-test rate\")\n parser.add_argument(\"--seed\", type=int, default=\"29\")\n\n # pre-processing\n parser.add_argument(\"--fft\", type=bool, default=False, help=\"FFT preprocessing\")\n parser.add_argument(\"--window\", type=int, default=128, help=\"time window, if not augment data, window=1024\")\n parser.add_argument(\"--normalization\", type=str, default=\"0-1\", choices=[\"None\", \"0-1\", \"mean-std\"], help=\"normalization option\")\n parser.add_argument(\"--savemodel\", type=bool, default=False, help=\"whether save pre-trained model in the classification task\")\n parser.add_argument(\"--pretrained\", type=bool, default=False, help=\"whether use pre-trained model in transfer learning tasks\")\n\n # backbone\n parser.add_argument(\"--backbone\", type=str, default=\"ResNet1D\", choices=[\"ResNet1D\", \"ResNet2D\", \"MLPNet\", \"CNN1D\"])\n # if backbone in (\"ResNet1D\", \"CNN1D\"), data shape: (batch size, 1, 1024)\n # elif backbone == \"ResNet2D\", data shape: (batch size, 3, 32, 32)\n # elif backbone == \"MLPNet\", data shape: (batch size, 1024)\n\n\n # optimization & training\n parser.add_argument(\"--num_workers\", type=int, default=0, help=\"the number of dataloader workers\")\n parser.add_argument(\"--batch_size\", type=int, default=256)\n parser.add_argument(\"--max_epoch\", type=int, default=100)\n parser.add_argument(\"--lr\", type=float, default=1e-3, help=\"learning rate\")\n parser.add_argument('--lr_scheduler', type=str, default='stepLR', choices=['step', 'exp', 'stepLR', 'fix'], help='the learning rate schedule')\n parser.add_argument('--gamma', type=float, default=0.8, help='learning rate scheduler parameter for step and exp')\n parser.add_argument('--steps', type=str, default='30, 120', help='the learning rate decay for step and stepLR')\n parser.add_argument(\"--optimizer\", type=str, default=\"adam\", choices=[\"adam\", \"sgd\"])\n parser.add_argument(\"--kernel\", type=str, default='Linear', choices=[\"Linear\", \"CORAL\"])\n\n\n args = parser.parse_args()\n return args\n\n# ===== Build Model =====\nclass FeatureNet(nn.Module):\n def __init__(self, args):\n super(FeatureNet, self).__init__()\n if args.backbone == \"ResNet1D\":\n self.feature_net = ResNet1D.resnet18()\n elif args.backbone == \"ResNet2D\":\n self.model_ft = models.resnet18(pretrained=True)\n self.bottleneck = nn.Sequential(nn.Linear(self.model_ft.fc.out_features, 512), nn.ReLU(), nn.Dropout(0.5))\n self.feature_net = nn.Sequential(self.model_ft, self.bottleneck)\n elif args.backbone == \"MLPNet\":\n if args.fft:\n self.feature_net = MLPNet.MLPNet(num_in=512)\n else:\n self.feature_net = MLPNet.MLPNet()\n elif args.backbone == \"CNN1D\":\n self.feature_net = CNN1D.CNN1D()\n else:\n raise Exception(\"model not implement\")\n\n def forward(self, x):\n logits = self.feature_net(x)\n\n return logits\n\nclass Classifier(nn.Module):\n def __init__(self, args, num_out=10):\n super(Classifier, self).__init__()\n if args.backbone in (\"ResNet1D\", \"ResNet2D\"):\n self.classifier = nn.Sequential(nn.Linear(512,num_out, nn.Dropout(0.5)))\n if args.backbone in (\"MLPNet\", \"CNN1D\"):\n self.classifier = nn.Sequential(nn.Linear(64,num_out, nn.Dropout(0.5)))\n\n def forward(self, logits):\n outputs = self.classifier(logits)\n\n return outputs\n\n# ===== Load Data =====\ndef loaddata(args):\n if args.source_dataname == \"CWRU\":\n source_data, source_label = CWRUloader(args, args.s_load, args.s_label_set)\n else:\n raise NotImplementedError(\"Source dataset {} not implemented.\".format(args.source_dataname))\n\n source_data, source_label = np.concatenate(source_data, axis=0), np.concatenate(source_label, axis=0)\n \n if args.target_dataname == \"CWRU\":\n target_data, target_label = CWRUloader(args, args.t_load, args.t_label_set)\n else:\n raise NotImplementedError(\"Target dataset {} not implemented.\".format(args.target_dataname))\n\n target_data, target_label = np.concatenate(target_data, axis=0), np.concatenate(target_label, axis=0)\n\n source_loader, _, _ = utils.DataSplite(args, source_data, source_label)\n target_trainloader, target_valloader, target_testloader = utils.DataSplite(args, target_data, target_label)\n \n return source_loader, target_trainloader, target_valloader, target_testloader\n\n# ===== Test the Model =====\ndef tester(featurenet, classifier, dataloader):\n featurenet.eval()\n classifier.eval()\n device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n correct_num, total_num = 0, 0\n for i, (x_batch, y_batch) in enumerate(dataloader):\n x_batch, y_batch = x_batch.to(device), y_batch.to(device)\n\n # compute model cotput and loss\n logtis_batch = featurenet(x_batch)\n output_batch = classifier(logtis_batch)\n\n pre = torch.max(output_batch.cpu(), 1)[1].numpy()\n y = y_batch.cpu().numpy()\n correct_num += (pre == y).sum()\n total_num += len(y)\n accuracy = (correct_num / total_num) * 100.0\n return accuracy\n\n\n# ===== Train the Model =====\ndef trainer(args):\n # Consider the gpu or cpu condition\n if torch.cuda.is_available():\n device = torch.device(\"cuda\")\n device_count = torch.cuda.device_count()\n logging.info('using {} gpus'.format(device_count))\n assert args.batch_size % device_count == 0, \"batch size should be divided by device count\"\n else:\n warnings.warn(\"gpu is not available\")\n device = torch.device(\"cpu\")\n device_count = 1\n logging.info('using {} cpu'.format(device_count))\n \n # load the dataset\n source_trainloader, target_trainloader, target_valloader, target_testloader = loaddata(args)\n\n # load the model\n featurenet = FeatureNet(args)\n classifier = Classifier(args, num_out=len(args.t_label_set))\n\n # load the checkpoint\n if args.pretrained:\n if args.backbone != \"ResNet2D\": # pretrained ResNet2D model is downloaded from torchvision module\n if not args.fft:\n path = \"./checkpoints/{}_checkpoint.tar\".format(args.backbone)\n else:\n path = \"./checkpoints/{}FFT_checkpoint.tar\".format(args.backbone)\n featurenet.load_state_dict(torch.load(path))\n\n parameter_list = [{\"params\": featurenet.parameters(), \"lr\": args.lr},\n {\"params\": classifier.parameters(), \"lr\": args.lr}]\n\n # Define optimizer and learning rate decay\n optimizer, lr_scheduler = utils.optimizer(args, parameter_list)\n\n # define loss function\n loss_cls = nn.CrossEntropyLoss()\n if args.kernel == \"Linear\":\n loss_dis = MMDLinear.MMDLinear\n elif args.kernel == \"CORAL\":\n loss_dis = CORAL.CORAL_loss\n else:\n raise NotImplemented(\"Kernel {} not implemented.\".format(args.kernel))\n\n featurenet.to(device)\n classifier.to(device)\n\n # train\n best_acc = 0.0\n meters = {\"acc_source_train\":[], \"acc_target_train\": [], \"acc_target_val\": []}\n\n for epoch in range(args.max_epoch):\n featurenet.train()\n classifier.train()\n with tqdm(total=len(target_trainloader), leave=False) as pbar:\n for i, ((x_s_batch, y_s_batch), (x_t_batch, y_t_batch)) in enumerate(zip(source_trainloader,target_trainloader)):\n\n if len(y_s_batch) != len(y_t_batch):\n break\n batch_num = x_s_batch.size(0)\n\n inputs = torch.cat((x_s_batch, x_t_batch), dim=0)\n\n # move to GPU if available\n inputs = inputs.to(device)\n s_labels = y_s_batch.to(device)\n t_labels = y_t_batch.to(device)\n\n # compute model cotput and loss\n logits = featurenet(inputs)\n outputs = classifier(logits)\n\n classification_loss = loss_cls(outputs.narrow(0, 0, s_labels.size(0)), s_labels.long())\n distance_loss = loss_dis(outputs.view(outputs.size(0),-1).narrow(0, 0, s_labels.size(0)),\\\n outputs.view(outputs.size(0),-1).narrow(0, s_labels.size(0), s_labels.size(0)))\n loss = classification_loss + distance_loss \n\n # clear previous gradients, compute gradients\n optimizer.zero_grad()\n loss.backward()\n\n # performs updates using calculated gradients\n optimizer.step()\n\n # evaluate\n # training accuracy\n acc_source_train = utils.accuracy(outputs.narrow(0, 0, batch_num), s_labels)\n acc_target_train = utils.accuracy(outputs.narrow(0, batch_num, batch_num), t_labels) \n\n pbar.update()\n \n # update lr\n if lr_scheduler is not None:\n lr_scheduler.step()\n\n val_acc = tester(featurenet, classifier, target_valloader)\n if val_acc > best_acc:\n best_acc = val_acc\n if args.savemodel:\n utils.save_model(featurenet, args)\n \n logging.info(\"Epoch: {:>3}/{}, loss_cls: {:.4f}, loss: {:.4f}, source_train_acc: {:>6.2f}%, target_train_acc: {:>6.2f}%, target_val_acc: {:>6.2f}%\".format(\\\n epoch+1, args.max_epoch, classification_loss, loss, acc_source_train, acc_target_train, val_acc))\n meters[\"acc_source_train\"].append(acc_source_train)\n meters[\"acc_target_train\"].append(acc_target_train)\n meters[\"acc_target_val\"].append(val_acc)\n\n logging.info(\"Best accuracy: {:.4f}\".format(best_acc))\n utils.save_log(meters, \"./logs/DDC_{}_{}_meters.pkl\".format(args.backbone, args.max_epoch))\n\n logging.info(\"=\"*15+\"Done!\"+\"=\"*15)\n\nif __name__ == \"__main__\":\n\n args = parse_args()\n\n # set the logger\n if not os.path.exists(\"./logs\"):\n os.makedirs(\"./logs\")\n setlogger(args.log_file)\n\n # save the args\n for k, v in args.__dict__.items():\n logging.info(\"{}: {}\".format(k, v))\n \n trainer(args)" }, { "path": "OSDABP.py", "content": "import argparse\nimport os\nimport numpy as np\nfrom Utils.logger import setlogger\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision.models as models\nfrom torch.autograd import Function\n\nfrom Backbone import ResNet1D, MLPNet, CNN1D\nfrom PreparData.CWRU import CWRUloader\nimport Utils.utils as utils\n\nfrom tqdm import *\nimport warnings\nimport logging\n\n# ===== Define argments =====\ndef parse_args():\n parser = argparse.ArgumentParser(description='Implementation of Deep Domain Confusion networks')\n\n # task setting\n parser.add_argument(\"--log_file\", type=str, default=\"./logs/OSDABP.log\", help=\"log file path\")\n\n # dataset information\n parser.add_argument(\"--datadir\", type=str, default=\"./datasets\", help=\"data directory\")\n parser.add_argument(\"--source_dataname\", type=str, default=\"CWRU\", choices=[\"CWRU\", \"PU\"], help=\"choice a dataset\")\n parser.add_argument(\"--target_dataname\", type=str, default=\"CWRU\", choices=[\"CWRU\", \"PU\"], help=\"choice a dataset\")\n parser.add_argument(\"--s_load\", type=int, default=3, help=\"source domain working condition\")\n parser.add_argument(\"--t_load\", type=int, default=2, help=\"target domain working condition\")\n parser.add_argument(\"--s_label_set\", type=list, default=[0,1,2,3,4,5], help=\"source domain label set\")\n parser.add_argument(\"--t_label_set\", type=list, default=[0,1,2,3,4,5,6,7,8,9], help=\"target domain label set\")\n parser.add_argument(\"--val_rat\", type=float, default=0.3, help=\"training-validation rate\")\n parser.add_argument(\"--test_rat\", type=float, default=0.5, help=\"validation-test rate\")\n parser.add_argument(\"--seed\", type=int, default=\"29\")\n\n # pre-processing\n parser.add_argument(\"--fft\", type=bool, default=False, help=\"FFT preprocessing\")\n parser.add_argument(\"--window\", type=int, default=128, help=\"time window, if not augment data, window=1024\")\n parser.add_argument(\"--normalization\", type=str, default=\"0-1\", choices=[\"None\", \"0-1\", \"mean-std\"], help=\"normalization option\")\n parser.add_argument(\"--savemodel\", type=bool, default=False, help=\"whether save pre-trained model in the classification task\")\n parser.add_argument(\"--pretrained\", type=bool, default=False, help=\"whether use pre-trained model in transfer learning tasks\")\n\n # backbone\n parser.add_argument(\"--backbone\", type=str, default=\"ResNet1D\", choices=[\"ResNet1D\", \"ResNet2D\", \"MLPNet\", \"CNN1D\"])\n # if backbone in (\"ResNet1D\", \"CNN1D\"), data shape: (batch size, 1, 1024)\n # elif backbone == \"ResNet2D\", data shape: (batch size, 3, 32, 32)\n # elif backbone == \"MLPNet\", data shape: (batch size, 1024)\n\n\n # optimization & training\n parser.add_argument(\"--num_workers\", type=int, default=0, help=\"the number of dataloader workers\")\n parser.add_argument(\"--batch_size\", type=int, default=256)\n parser.add_argument(\"--max_epoch\", type=int, default=100)\n parser.add_argument(\"--lr\", type=float, default=1e-3, help=\"learning rate\")\n parser.add_argument('--lr_scheduler', type=str, default='stepLR', choices=['step', 'exp', 'stepLR', 'fix'], help='the learning rate schedule')\n parser.add_argument('--gamma', type=float, default=0.8, help='learning rate scheduler parameter for step and exp')\n parser.add_argument('--steps', type=str, default='30, 120', help='the learning rate decay for step and stepLR')\n parser.add_argument(\"--optimizer\", type=str, default=\"adam\", choices=[\"adam\", \"sgd\"])\n\n\n args = parser.parse_args()\n return args\n\n# ===== Build Model =====\nclass GradReverse(Function):\n @staticmethod\n def forward(ctx, x):\n return x.view_as(x)\n\n @staticmethod\n def backward(ctx, grad_output):\n output = grad_output.neg()\n return output, None\n\ndef grad_reverse(x):\n return GradReverse.apply(x)\n\n# define the model\nclass FeatureNet(nn.Module):\n def __init__(self, args):\n super(FeatureNet, self).__init__()\n if args.backbone == \"ResNet1D\":\n self.feature_net = ResNet1D.resnet18()\n elif args.backbone == \"ResNet2D\":\n self.model_ft = models.resnet18(pretrained=True)\n self.bottleneck = nn.Sequential(nn.Linear(self.model_ft.fc.out_features, 512), nn.ReLU(), nn.Dropout(0.2))\n self.feature_net = nn.Sequential(self.model_ft, self.bottleneck)\n elif args.backbone == \"MLPNet\":\n self.feature_net = MLPNet.MLPNet()\n elif args.backbone == \"CNN1D\":\n self.feature_net = CNN1D.CNN1D()\n else:\n raise Exception(\"model not implement\")\n\n def forward(self, x):\n logits = self.feature_net(x)\n\n return logits\n\nclass Classifier(nn.Module):\n def __init__(self, args, num_out=10):\n super(Classifier, self).__init__()\n if args.backbone in (\"ResNet1D\", \"ResNet2D\"):\n self.classifier = nn.Sequential(nn.Linear(512,num_out, nn.Dropout(0.5)))\n if args.backbone in (\"MLPNet\", \"CNN1D\"):\n self.classifier = nn.Sequential(nn.Linear(64,num_out, nn.Dropout(0.5)))\n\n def forward(self, logits, reverse = False):\n if reverse:\n logits = grad_reverse(logits)\n outputs = self.classifier(logits)\n\n return outputs\n\n# ===== Load Data =====\ndef loaddata(args):\n if args.source_dataname == \"CWRU\":\n source_data, source_label = CWRUloader(args, args.s_load, args.s_label_set)\n\n source_data, source_label = np.concatenate(source_data, axis=0), np.concatenate(source_label, axis=0)\n \n if args.target_dataname == \"CWRU\":\n target_data, target_label = CWRUloader(args, args.t_load, args.t_label_set)\n\n target_data, target_label = np.concatenate(target_data, axis=0), np.concatenate(target_label, axis=0)\n\n source_loader, _, _ = utils.DataSplite(args, source_data, source_label)\n target_trainloader, target_valloader, target_testloader = utils.DataSplite(args, target_data, target_label)\n \n return source_loader, target_trainloader, target_valloader, target_testloader\n\n# ===== Define Loss Function =====\ndef bce_loss(output, target):\n output_neg = 1 - output\n target_neg = 1 - target\n result = torch.mean(target * torch.log(output + 1e-6))\n result += torch.mean(target_neg * torch.log(output_neg + 1e-6))\n return -torch.mean(result)\n\n# ===== Test the Model =====\ndef tester(featurenet, classifier, dataloader):\n featurenet.eval()\n classifier.eval()\n device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n correct_num, total_num = 0, 0\n for i, (x_batch, y_batch) in enumerate(dataloader):\n x_batch, y_batch = x_batch.to(device), y_batch.to(device)\n\n # compute model cotput and loss\n logtis_batch = featurenet(x_batch)\n output_batch = classifier(logtis_batch)\n\n pre = torch.max(output_batch.cpu(), 1)[1].numpy()\n y = y_batch.cpu().numpy()\n correct_num += (pre == y).sum()\n total_num += len(y)\n accuracy = (correct_num / total_num) * 100.0\n return accuracy\n\n\n# ===== Train the Model =====\ndef trainer(args):\n # Consider the gpu or cpu condition\n if torch.cuda.is_available():\n device = torch.device(\"cuda\")\n device_count = torch.cuda.device_count()\n logging.info('using {} gpus'.format(device_count))\n assert args.batch_size % device_count == 0, \"batch size should be divided by device count\"\n else:\n warnings.warn(\"gpu is not available\")\n device = torch.device(\"cpu\")\n device_count = 1\n logging.info('using {} cpu'.format(device_count))\n \n # load the dataset\n source_trainloader, target_trainloader, target_valloader, target_testloader = loaddata(args)\n\n num_out = len(args.s_label_set)+1\n\n # load the model\n featurenet = FeatureNet(args)\n classifier = Classifier(args, num_out=num_out)\n\n # load the checkpoint\n if args.pretrained:\n if args.backbone != \"ResNet2D\": # pretrained ResNet2D model is downloaded from torchvision module\n if not args.fft:\n path = \"./checkpoints/{}_checkpoint.tar\".format(args.backbone)\n else:\n path = \"./checkpoints/{}FFT_checkpoint.tar\".format(args.backbone)\n featurenet.load_state_dict(torch.load(path))\n\n parameter_list = [{\"params\": featurenet.parameters(), \"lr\": args.lr},\n {\"params\": classifier.parameters(), \"lr\": args.lr}]\n\n # Define optimizer and learning rate decay\n optimizer, lr_scheduler = utils.optimizer(args, parameter_list)\n\n ## define loss function\n criterion = nn.CrossEntropyLoss()\n\n featurenet.to(device)\n classifier.to(device)\n\n # train\n best_acc = 0.0\n meters = {\"acc_source_train\":[], \"acc_target_train\": [], \"acc_target_val\": []}\n\n for epoch in range(args.max_epoch):\n featurenet.train()\n classifier.train()\n with tqdm(total=len(target_trainloader), leave=False) as pbar:\n for i, ((x_s_batch, y_s_batch), (x_t_batch, y_t_batch)) in enumerate(zip(source_trainloader,target_trainloader)):\n\n if len(y_s_batch) != len(y_t_batch):\n break\n batch_num = x_s_batch.size(0)\n\n target_funk = torch.FloatTensor(batch_num, 2).fill_(0.5).cuda()\n\n # clear previous gradients, compute gradients\n optimizer.zero_grad()\n\n # move to GPU if available\n x_s = x_s_batch.to(device)\n x_t = x_t_batch.to(device)\n y_s = y_s_batch.to(device)\n y_t = y_t_batch.to(device)\n\n # compute model output and loss\n # source data\n logits_s = featurenet(x_s)\n outputs_s = classifier(logits_s)\n loss_s = criterion(outputs_s, y_s.long())\n loss_s.backward()\n\n # target data\n logits_t = featurenet(x_t)\n outputs_t = classifier(logits_t, reverse=True)\n outputs_t = F.softmax(outputs_t)\n prob1 = torch.sum(outputs_t[:, :num_out-1], 1).view(-1, 1)\n prob2 = outputs_t[:, num_out-1].contiguous().view(-1, 1)\n prob = torch.cat([prob1, prob2], 1)\n\n loss_t = bce_loss(prob, target_funk)\n loss_t.backward()\n\n # performs updates using calculated gradients\n optimizer.step()\n # clear previous gradients\n optimizer.zero_grad() \n\n pbar.update()\n \n # update lr\n if lr_scheduler is not None:\n lr_scheduler.step()\n\n # validation\n featurenet.eval()\n classifier.eval()\n correct_num = 0\n val_num = 0\n per_class_num = np.zeros((num_out))\n per_class_correct = np.zeros((num_out)).astype(np.float32)\n for step, (x_val_batch, y_val_batch) in enumerate(target_valloader):\n # move to GPU if available\n x_val= x_val_batch.to(device)\n y_val = y_val_batch.to(device)\n batch_size_val = y_val.data.size()[0]\n\n logits_val = featurenet(x_val)\n outputs_val = classifier(logits_val)\n\n pre = torch.max(outputs_val.cpu(), 1)[1].numpy()\n y_val = y_val.cpu().numpy()\n\n correct_num += (pre == y_val).sum() # the number of correct preditions per batch\n val_num += batch_size_val # the number of predictions per batch\n\n for i in range(num_out):\n if i < num_out -1:\n index = np.where(y_val == i) # known classes\n else:\n index = np.where(y_val >= i) # unknown classes \n # Thanks to @Wang-Dongdong for reporting the bug\n correct_ind = np.where(pre[index[0]]==i)\n per_class_correct[i] += float(len(correct_ind[0]))\n per_class_num[i] += float(len(index[0]))\n\n per_class_acc = (per_class_correct / per_class_num) * 100.0\n known_acc = (per_class_correct[:-1].sum() / per_class_num[:-1].sum()) * 100.0\n all_acc = (correct_num / val_num) * 100.0\n\n if all_acc > best_acc:\n best_acc = all_acc\n logging.info(\"Epoch: {:>3}/{}, loss_s: {:.4f}, loss_t: {:.4f}, all_acc: {:>6.2f}, known_acc: {:>6.2f}%\".format(\\\n epoch+1, args.max_epoch, loss_s, loss_t, all_acc, known_acc))\n\n logging.info(\"Best all accuracy: {:.4f}\".format(best_acc))\n\n logging.info(\"=\"*10+\"Done!\"+\"=\"*10)\n\nif __name__ == \"__main__\":\n\n args = parse_args()\n\n # set the logger\n if not os.path.exists(\"./logs\"):\n os.makedirs(\"./logs\")\n setlogger(args.log_file)\n\n # save the args\n for k, v in args.__dict__.items():\n logging.info(\"{}: {}\".format(k, v))\n \n trainer(args)" }, { "path": "PreparData/CWRU.py", "content": "\"\"\"\n@Author: Xiaohan Chen\n@Email: cxh_bb@outlook.com\n\"\"\"\n\nimport numpy as np\nfrom scipy.io import loadmat\nfrom PreparData.preprocess import transformation\n\n# datanames in every working conditions\ndataname_dict= {0:[97, 109, 122, 135, 174, 189, 201, 213, 226, 238], # 1797rpm\n 1:[98, 110, 123, 136, 175, 190, 202, 214, 227, 239], # 1772rpm\n 2:[99, 111, 124, 137, 176, 191, 203, 215, 228, 240], # 1750rpm\n 3:[100,112, 125, 138, 177, 192, 204, 217, 229, 241]} # 1730rpm\n\naxis = \"_DE_time\"\ndata_length = 1024\n\n\ndef CWRU(datadir, load, labels, window, normalization, backbone, fft):\n \"\"\"\n loading the hole dataset\n \"\"\"\n path = datadir + \"/CWRU/\" + \"Drive_end_\" + str(load) + \"/\"\n dataset = {label: [] for label in labels}\n for label in labels:\n fault_type = dataname_dict[load][label]\n if fault_type < 100:\n realaxis = \"X0\" + str(fault_type) + axis\n else:\n realaxis = \"X\" + str(fault_type) + axis\n mat_data = loadmat(path+str(fault_type)+\".mat\")[realaxis]\n start, end = 0, data_length\n\n # set the endpoint of data sequence\n endpoint = mat_data.shape[0]\n\n # split the data and transformation\n while end < endpoint:\n sub_data = mat_data[start : end].reshape(-1,)\n\n sub_data = transformation(sub_data, fft, normalization, backbone)\n\n dataset[label].append(sub_data)\n start += window\n end += window\n \n dataset[label] = np.array(dataset[label], dtype=\"float32\")\n\n return dataset\n\ndef CWRUloader(args, load, label_set, number=\"all\"):\n \"\"\"\n args: arguments\n number: the numbers of training samples, \"all\" or specific numbers (string type)\n \"\"\"\n dataset = CWRU(args.datadir, load, label_set, args.window, args.normalization, args.backbone, args.fft)\n\n DATA, LABEL = [], []\n\n if number == \"all\":\n counter = []\n for key in dataset.keys():\n counter.append(dataset[key].shape[0])\n datan = min(counter) # choosing the min value as the sample size per class\n for key in dataset.keys():\n LABEL.append(np.tile(key, datan))\n DATA.append(dataset[key][:datan])\n else:\n datan = int(number)\n for key in dataset.keys():\n LABEL.append(np.tile(key, datan))\n DATA.append(dataset[key][:datan])\n \n DATA, LABEL = np.array(DATA, dtype=\"float32\"), np.array(LABEL, dtype=\"int32\")\n\n return DATA, LABEL" }, { "path": "PreparData/__init__.py", "content": "" }, { "path": "PreparData/preprocess.py", "content": "import numpy as np\n\ndef transformation(sub_data, fft, normalization, backbone):\n\n if fft:\n sub_data = np.fft.fft(sub_data)\n sub_data = np.abs(sub_data) / len(sub_data)\n sub_data = sub_data[:int(sub_data.shape[0] / 2)].reshape(-1,) \n\n if normalization == \"0-1\":\n sub_data = (sub_data - sub_data.min()) / (sub_data.max() - sub_data.min())\n elif normalization == \"mean-std\":\n sub_data = (sub_data - sub_data.mean()) / sub_data.std()\n\n if backbone in (\"ResNet1D\", \"CNN1D\"):\n sub_data = sub_data[np.newaxis, :]\n elif backbone == \"ResNet2D\":\n n = int(np.sqrt(sub_data.shape[0]))\n if fft:\n sub_data = sub_data[:n*n]\n sub_data = np.reshape(sub_data, (n, n))\n sub_data = sub_data[np.newaxis, :]\n sub_data = np.concatenate((sub_data, sub_data, sub_data), axis=0)\n\n return sub_data" }, { "path": "README.md", "content": "# Deep transfer learing for fault diagnosis\n\n## :book: 1. Introduction\nThis repository contains popular deep transfer learning algorithms implemented via PyTorch for cross-load fault diagnosis transfer tasks, including: \n\n- [x] General supervised learning classification task: traing and test apply the same machines, working conditions and faults.\n\n- [x] *domain adaptation*: the distribution of the source domain data may be different from the target domain data, but the label set of the target domain is the same as the source domain, i.e., $\\mathcal{D} _{s}=(X_s,Y_s)$, $\\mathcal{D} _{t}=(X_t,Y_t)$, $X_s \\ne X_t$, $Y_s = Y_t$.\n - [x] **DDC**: Deep Domain Confusion [[arXiv 2014]](https://arxiv.org/pdf/1412.3474.pdf)\n - [x] **Deep CORAL**: Correlation Alignment for Deep Domain Adaptation [[ECCV 2016]](https://arxiv.org/abs/1607.01719)\n - [x] **DANN**: Unsupervised Domain Adaptation by Backpropagation [[ICML 2015]](http://proceedings.mlr.press/v37/ganin15.pdf)\n - [ ] TODO\n\n- [x] *Open-set domain adaptation*: the distribution of the source domain data may be different from the target domain data. What's more, the target label set contains unknown categories, i.e., $\\mathcal{D} _{s}=(X_s,Y_s)$, $\\mathcal{D} _{t}=(X_t,Y_t)$, $X_s \\ne X_t$, $Y_s \\in Y_t$. We refer to their common categories $\\mathcal{Y}_s\\cap \\mathcal{Y}_t$ as the *known classes*, and $\\mathcal{Y}_s\\setminus \\mathcal{Y}_t$ (or $\\mathcal{Y}_t\\setminus \\mathcal{Y}_s$) in the target domain as the *unknown class*.\n - [x] **OSDABP**: Open Set Domain Adaptation by Backpropagation [[ECCV 2018]](http://openaccess.thecvf.com/content_ECCV_2018/papers/Kuniaki_Saito_Adversarial_Open_Set_ECCV_2018_paper.pdf)\n - [ ] TODO\n\n> **Few-shot** learning-based bearing fault diagnosis methods please see: https://github.com/Xiaohan-Chen/few-shot-fault-diagnosis\n\n## :balloon: 2. Citation\n\nFor further introductions to transfer learning in bearing fault diagnosis, please read our [paper](https://ieeexplore.ieee.org/document/10042467). And if you find this repository useful and use it in your works, please cite our paper, thank you~:\n```\n@ARTICLE{10042467,\n author={Chen, Xiaohan and Yang, Rui and Xue, Yihao and Huang, Mengjie and Ferrero, Roberto and Wang, Zidong},\n journal={IEEE Transactions on Instrumentation and Measurement}, \n title={Deep Transfer Learning for Bearing Fault Diagnosis: A Systematic Review Since 2016}, \n year={2023},\n volume={72},\n number={},\n pages={1-21},\n doi={10.1109/TIM.2023.3244237}}\n```\n\n---\n## :wrench: 3. Requirements\n- python 3.9.12\n- Numpy 1.23.1\n- pytorch 1.12.0\n- scikit-learn 1.1.1\n- torchvision 0.13.0\n\n---\n## :handbag: 4. Dataset\nDownload the bearing dataset from [CWRU Bearing Dataset Center](https://engineering.case.edu/bearingdatacenter/48k-drive-end-bearing-fault-data) and place the `.mat` files in the `./datasets` folder according to the following structure:\n```\ndatasets/\n └── CWRU/\n ├── Drive_end_0/\n │ └── 97.mat 109.mat 122.mat 135.mat 174.mat 189.mat 201.mat 213.mat 226.mat 238.mat\n ├── Drive_end_1/\n │ └── 98.mat 110.mat 123.mat 136.mat 175.mat 190.mat 202.mat 214.mat 227.mat 239.mat\n ├── Drive_end_2/\n │ └── 99.mat 111.mat 124.mat 137.mat 176.mat 191.mat 203.mat 215.mat 228.mat 240.mat\n └── Drive_end_3/\n └── 100.mat 112.mat 125.mat 138.mat 177.mat 192.mat 204.mat 217.mat 229.mat 241.mat\n```\n\n---\n## :pencil: 5. Usage\n> **NOTE**: When using pre-trained models to initialise the backbone and classifier in transfer learning tasks, run classification tasks first to generate corresponding checkpoints.\n\nFour typical neural networks are implemented in this repository, including MLP, 1D CNN, 1D ResNet18, and 2D ResNet18(torchvision package). More details can be found in the `./Backbone` folder.\n\n**General Supervised Learning Classification:**\n- Train and test the model on the same machines, working conditions and faults. Use the following commands:\n```python\npython3 classification.py --datadir './datasets' --max_epoch 100\n```\n\n**Transfer Learning:**\n- If using the DDC transfer learning method, use the following commands:\n```python\npython3 DDC.py --datadir './datasets' --backbone \"CNN1D\" --pretrained False --kernel 'Linear'\n```\n- If using the DeepCORAL transfer learning method, use the following commands:\n```python\npython3 DDC.py --datadir './datasets' --backbone \"CNN1D\" --pretrained False --kernel 'CORAL'\n```\n- If using the DANN transfer learning method, use following commands:\n```python\npython3 DANN.py --backbone \"CNN1D\"\n```\n**Open Set Domain Adaptation:**\n- The target domain contains unknow classes, use the following commands:\n```python\npython3 OSDABP.py\n```\n---\n## :flashlight: 6. Results\n> The following results do not represent the best results.\n\n**General Classification task:** \nDataset: CWRU \nLoad: 3 \nLabel set: [0,1,2,3,4,5,6,7,8,9] \n\n| | MLPNet | CNN1D | ResNet1D | ResNet2D |\n| :---------------: | :----: | :---: | :------: | :------: |\n| acc (time domain) | 93.95 | 97.70 | 99.58 | 98.02 |\n| acc (freq domain) | 99.95 | 99.44 | 100.0 | 99.96 |\n\n**Transfer Learning:** \nDataset: CWRU \nSource load: 3 \nTarget Load: 2 \nLabel set: [0,1,2,3,4,5,6,7,8,9] \nPre-trained model: True \n\nTime domain: \n| | MLPNet | CNN1D | ResNet1D | ResNet2D |\n| :-----------------: | :----: | :---: | :------: | :------: |\n| DDC (linear kernel) | 75.47 | 85.53 | 91.79 | 91.32 |\n| DeepCORAL | 82.33 | 88.23 | 93.88 | 90.84 |\n| DANN | 87.68 | 94.77 | 98.88 | 93.95 |\n\nFrequency domain\n| | MLPNet | CNN1D | ResNet1D | ResNet2D |\n| :-------: | :----: | :---: | :------: | :------: |\n| DeepCORAL | 98.65 | 98.22 | 99.75 | 99.31 |\n| DANN | 99.38 | 98.74 | 99.89 | 99.47 |\n\n**Open Set Domain Adaptation** \n- *OSDABP*\nDataset: CWRU \nSource load: 3 \nTarget Load: 2 \nSource label set: [0,1,2,3,4,5] \nTarget label set: [0,1,2,3,4,5,6,7,8,9] \nPre-trained model: True \n\n| Label | 0 | 1 | 2 | 3 | 4 | 5 | unk | All | Only known |\n| :------: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | ----- | ---------- |\n| MLPNet | 99.83 | 95.96 | 59.76 | 76.10 | 19.85 | 96.58 | 59.21 | 70.21 | 75.99 |\n| CNN1D | 100.0 | 94.95 | 94.47 | 99.08 | 47.31 | 74.32 | 26.36 | 61.75 | 85.35 |\n| ResNet1D | 100.0 | 100.0 | 80.14 | 100.0 | 43.32 | 93.49 | 45.22 | 70.04 | 86.58 |\n| ResNet2D | 100.0 | 100.0 | 94.82 | 100.0 | 18.55 | 98.12 | 53.42 | 72.95 | 85.96 |\n\n\n---\n## :camping: 7. See also\n- Multi-scale CNN and LSTM bearing fault diagnosis [[paper](https://link.springer.com/article/10.1007/s10845-020-01600-2)][[GitHub](https://github.com/Xiaohan-Chen/baer_fault_diagnosis)]\n- TFPred self-supervised learning for few labeled fault diagnosis [[Paper](https://www.sciencedirect.com/science/article/pii/S0967066124000601)][[GitHub](https://github.com/Xiaohan-Chen/TFPred)]\n\n---\n## :globe_with_meridians: 8. Acknowledgement\n\n```\n@article{zhao2021applications,\n title={Applications of Unsupervised Deep Transfer Learning to Intelligent Fault Diagnosis: A Survey and Comparative Study},\n author={Zhibin Zhao and Qiyang Zhang and Xiaolei Yu and Chuang Sun and Shibin Wang and Ruqiang Yan and Xuefeng Chen},\n journal={IEEE Transactions on Instrumentation and Measurement},\n year={2021}\n}\n```\n\nI would like to thank the following person for contributing to this repository: [@Wang-Dongdong](https://github.com/Wang-Dongdong),[@zhuting233](https://github.com/zhuting233)" }, { "path": "Utils/__init__.py", "content": "" }, { "path": "Utils/logger.py", "content": "import logging\n\ndef setlogger(path):\n\n logger = logging.getLogger()\n logger.setLevel(logging.INFO)\n formatter = logging.Formatter(\"%(asctime)s | %(message)s\", datefmt=\"%Y-%m-%d %H:%M:%S\")\n\n consoleHandler = logging.StreamHandler()\n fileHandler = logging.FileHandler(filename=path)\n\n consoleHandler.setFormatter(formatter)\n fileHandler.setFormatter(formatter)\n\n logger.addHandler(consoleHandler)\n logger.addHandler(fileHandler)" }, { "path": "Utils/utils.py", "content": "import logging\nimport torch\nimport pickle\nimport torch.optim as optim\nfrom sklearn.model_selection import train_test_split\n\ndef accuracy(outputs, labels):\n \"\"\"\n Compute the accuracy\n outputs, labels: (tensor)\n return: (float) accuracy in [0, 100]\n \"\"\"\n pre = torch.max(outputs.cpu(), 1)[1].numpy()\n y = labels.data.cpu().numpy()\n acc = ((pre == y).sum() / len(y)) * 100\n return acc\n\ndef save_log(obj, path):\n with open(path, \"wb\") as f:\n pickle.dump(obj, f)\n\ndef read_pkl(path):\n with open(path, 'rb') as f:\n data = pickle.load(f)\n return data\n\ndef save_model(model, args):\n if not args.fft:\n torch.save(model.state_dict(), \"./checkpoints/{}_checkpoint.tar\".format(args.backbone))\n else:\n torch.save(model.state_dict(), \"./checkpoints/{}FFT_checkpoint.tar\".format(args.backbone))\n\ndef DataSplite(args, data, label):\n \"\"\"\n split the data and lebel and transform the narray type to tensor type\n \"\"\"\n data_train, data_val, label_train, label_val = train_test_split(data, label, test_size = args.val_rat, random_state=args.seed)\n data_val, data_test, label_val, label_test = train_test_split(data_val, label_val, test_size= args.test_rat)\n\n # numpy to tensor\n data_train = torch.from_numpy(data_train).float()\n data_val = torch.from_numpy(data_val).float()\n data_test = torch.from_numpy(data_test).float()\n label_train = torch.from_numpy(label_train).float()\n label_val = torch.from_numpy(label_val).float()\n label_test = torch.from_numpy(label_test).float()\n\n # logging the data shape\n logging.info(\"training data/label shape: {},{}\".format(data_train.size(), label_train.size()))\n logging.info(\"validation data/label shape: {},{}\".format(data_val.size(), label_val.size()))\n logging.info(\"test data/label shape: {},{}\".format(data_test.size(), label_test.size()))\n\n # build the dataloader\n train = torch.utils.data.TensorDataset(data_train, label_train)\n val = torch.utils.data.TensorDataset(data_val, label_val)\n test = torch.utils.data.TensorDataset(data_test, label_test)\n\n train_loader = torch.utils.data.DataLoader(train, batch_size=args.batch_size, \\\n shuffle=True, num_workers=args.num_workers)\n\n val_loader = torch.utils.data.DataLoader(val, batch_size=args.batch_size, \\\n shuffle=False, num_workers=args.num_workers)\n test_loader = torch.utils.data.DataLoader(test, batch_size=args.batch_size, \\\n shuffle=False, num_workers=args.num_workers)\n\n return train_loader, val_loader, test_loader\n\ndef optimizer(args, parameter_list):\n # define optimizer\n if args.optimizer == \"sgd\":\n optimizer = optim.SGD(parameter_list, lr=args.lr, momentum=0.9, weight_decay=5e-4)\n elif args.optimizer == \"adam\":\n optimizer = optim.Adam(parameter_list, lr=args.lr)\n else:\n raise Exception(\"optimizer not implement\")\n\n # Define the learning rate decay\n if args.lr_scheduler == 'step':\n steps = [int(step) for step in args.steps.split(',')]\n lr_scheduler = optim.lr_scheduler.MultiStepLR(optimizer, steps, gamma=args.gamma)\n elif args.lr_scheduler == 'exp':\n lr_scheduler = optim.lr_scheduler.ExponentialLR(optimizer, args.gamma)\n elif args.lr_scheduler == 'stepLR':\n steps = int(args.steps.split(\",\")[0])\n lr_scheduler = optim.lr_scheduler.StepLR(optimizer, steps, args.gamma)\n elif args.lr_scheduler == 'cos':\n lr_scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, 20, 0)\n elif args.lr_scheduler == 'fix':\n lr_scheduler = None\n else:\n raise Exception(\"lr schedule not implement\")\n\n return optimizer, lr_scheduler\n" }, { "path": "classification.py", "content": "import argparse\nimport os\nimport numpy as np\nfrom Utils.logger import setlogger\n\nfrom turtle import forward\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torchvision.models as models\n\nfrom Backbone import ResNet1D, MLPNet, CNN1D\nfrom PreparData.CWRU import CWRUloader\nimport Utils.utils as utils\n\nfrom tqdm import *\nimport warnings\nimport logging\n\n# ===== Define argments =====\ndef parse_args():\n parser = argparse.ArgumentParser(description='classification task')\n\n # task setting\n parser.add_argument(\"--log_file\", type=str, default=\"./logs/classification.log\", help=\"log file path\")\n\n # dataset information\n parser.add_argument(\"--datadir\", type=str, default=\"./datasets\", help=\"data directory\")\n parser.add_argument(\"--load\", type=int, default=3, help=\"working condition\")\n parser.add_argument(\"--label_set\", type=list, default=[0,1,2,3,4,5,6,7,8,9], help=\"label set\")\n parser.add_argument(\"--val_rat\", type=float, default=0.3, help=\"training-validation rate\")\n parser.add_argument(\"--test_rat\", type=float, default=0.5, help=\"validation-test rate\")\n parser.add_argument(\"--seed\", type=int, default=\"29\")\n\n # pre-processing\n parser.add_argument(\"--fft\", type=bool, default=False, help=\"FFT preprocessing\")\n parser.add_argument(\"--window\", type=int, default=128, help=\"time window, if not augment data, window=1024\")\n parser.add_argument(\"--normalization\", type=str, default=\"0-1\", choices=[\"None\", \"0-1\", \"mean-std\"], help=\"normalization option\")\n parser.add_argument(\"--savemodel\", type=bool, default=False, help=\"whether save pre-trained model in the classification task\")\n parser.add_argument(\"--pretrained\", type=bool, default=False, help=\"whether use pre-trained model in transfer learning tasks\")\n\n # backbone\n parser.add_argument(\"--backbone\", type=str, default=\"ResNet1D\", choices=[\"ResNet1D\", \"ResNet2D\", \"MLPNet\", \"CNN1D\"])\n # if backbone in (\"ResNet1D\", \"CNN1D\"), data shape: (batch size, 1, 1024)\n # elif backbone == \"ResNet2D\", data shape: (batch size, 3, 32, 32)\n # elif backbone == \"MLPNet\", data shape: (batch size, 1024)\n\n\n # optimization & training\n parser.add_argument(\"--num_workers\", type=int, default=0, help=\"the number of dataloader workers\")\n parser.add_argument(\"--batch_size\", type=int, default=256)\n parser.add_argument(\"--max_epoch\", type=int, default=100)\n parser.add_argument(\"--lr\", type=float, default=1e-3, help=\"learning rate\")\n parser.add_argument('--lr_scheduler', type=str, default='stepLR', choices=['step', 'exp', 'stepLR', 'fix'], help='the learning rate schedule')\n parser.add_argument('--gamma', type=float, default=0.8, help='learning rate scheduler parameter for step and exp')\n parser.add_argument('--steps', type=str, default='30, 120', help='the learning rate decay for step and stepLR')\n parser.add_argument(\"--optimizer\", type=str, default=\"adam\", choices=[\"adam\", \"sgd\"])\n\n\n args = parser.parse_args()\n return args\n\n# ===== Build Model =====\nclass FeatureNet(nn.Module):\n def __init__(self, args):\n super(FeatureNet, self).__init__()\n if args.backbone == \"ResNet1D\":\n self.feature_net = ResNet1D.resnet18()\n elif args.backbone == \"ResNet2D\":\n self.model_ft = models.resnet18(pretrained=True)\n self.bottleneck = nn.Sequential(nn.Linear(self.model_ft.fc.out_features, 512), nn.ReLU(), nn.Dropout(0.5))\n self.feature_net = nn.Sequential(self.model_ft, self.bottleneck)\n elif args.backbone == \"MLPNet\":\n if args.fft:\n self.feature_net = MLPNet.MLPNet(num_in=512)\n else:\n self.feature_net = MLPNet.MLPNet()\n elif args.backbone == \"CNN1D\":\n self.feature_net = CNN1D.CNN1D()\n else:\n raise Exception(\"model not implement\")\n\n def forward(self, x):\n logits = self.feature_net(x)\n\n return logits\n\nclass Classifier(nn.Module):\n def __init__(self, args, num_out=10):\n super(Classifier, self).__init__()\n if args.backbone in (\"ResNet1D\", \"ResNet2D\"):\n self.classifier = nn.Sequential(nn.Linear(512,num_out, nn.Dropout(0.5)))\n if args.backbone in (\"MLPNet\", \"CNN1D\"):\n self.classifier = nn.Sequential(nn.Linear(64,num_out, nn.Dropout(0.5)))\n\n def forward(self, logits):\n outputs = self.classifier(logits)\n\n return outputs\n\n# ===== Load Data =====\ndef loaddata(args):\n data, label = CWRUloader(args, args.load, args.label_set)\n data, label = np.concatenate(data, axis=0), np.concatenate(label, axis=0)\n \n train_loader, val_loader, test_laoder = utils.DataSplite(args, data, label)\n \n return train_loader, val_loader, test_laoder\n\n# ===== Test the Model =====\ndef tester(featurenet, classifier, dataloader):\n featurenet.eval()\n classifier.eval()\n device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n correct_num, total_num = 0, 0\n for i, (x_batch, y_batch) in enumerate(dataloader):\n x_batch, y_batch = x_batch.to(device), y_batch.to(device)\n\n # compute model cotput and loss\n logtis_batch = featurenet(x_batch)\n output_batch = classifier(logtis_batch)\n\n pre = torch.max(output_batch.cpu(), 1)[1].numpy()\n y = y_batch.cpu().numpy()\n correct_num += (pre == y).sum()\n total_num += len(y)\n accuracy = (correct_num / total_num) * 100.0\n return accuracy\n\n\n# ===== Train the Model =====\ndef trainer(args):\n # Consider the gpu or cpu condition\n if torch.cuda.is_available():\n device = torch.device(\"cuda\")\n device_count = torch.cuda.device_count()\n logging.info('using {} gpus'.format(device_count))\n assert args.batch_size % device_count == 0, \"batch size should be divided by device count\"\n else:\n warnings.warn(\"gpu is not available\")\n device = torch.device(\"cpu\")\n device_count = 1\n logging.info('using {} cpu'.format(device_count))\n \n # load the dataset\n trainloader, valloader, testloader = loaddata(args)\n\n # load the model\n featurenet = FeatureNet(args)\n classifier = Classifier(args, num_out=len(args.label_set))\n\n parameter_list = [{\"params\": featurenet.parameters(), \"lr\": args.lr},\n {\"params\": classifier.parameters(), \"lr\": args.lr}]\n\n # Define optimizer and learning rate decay\n optimizer, lr_scheduler = utils.optimizer(args, parameter_list)\n\n # define loss function\n loss_fn = nn.CrossEntropyLoss()\n featurenet.to(device)\n classifier.to(device)\n # train\n best_acc = 0.0\n meters = {\"acc_train\": [], \"acc_val\": []}\n\n for epoch in range(args.max_epoch):\n featurenet.train()\n classifier.train()\n with tqdm(total=len(trainloader), leave=False) as pbar:\n for i, (x_batch, y_batch) in enumerate(trainloader):\n # move to GPU if available\n x_batch, y_batch = x_batch.to(device), y_batch.to(device)\n\n # compute model cotput and loss\n logtis_batch = featurenet(x_batch)\n output_batch = classifier(logtis_batch)\n loss = loss_fn(output_batch, y_batch.long())\n\n # clear previous gradients, compute gradients\n optimizer.zero_grad()\n loss.backward()\n\n # performs updates using calculated gradients\n optimizer.step()\n\n # evaluate\n # training accuracy\n train_acc = utils.accuracy(output_batch, y_batch)\n\n pbar.update()\n \n # update lr\n if lr_scheduler is not None:\n lr_scheduler.step()\n\n val_acc = tester(featurenet, classifier, valloader)\n if val_acc > best_acc:\n best_acc = val_acc\n if args.savemodel:\n utils.save_model(featurenet, args)\n \n logging.info(\"Epoch: {:>3}/{}, loss: {:.4f}, train_acc: {:>6.2f}%, val_acc: {:>6.2f}%\".format(\\\n epoch+1, args.max_epoch, loss, train_acc, val_acc))\n meters[\"acc_train\"].append(train_acc)\n meters[\"acc_val\"].append(val_acc)\n\n logging.info(\"Best accuracy: {:.4f}%\".format(best_acc))\n utils.save_log(meters, \"./logs/cls_{}_{}_meters.pkl\".format(args.backbone, args.max_epoch))\n\n logging.info(\"=\"*15+\"Done!\"+\"=\"*15)\n\nif __name__ == \"__main__\":\n\n args = parse_args()\n\n # set the logger\n if not os.path.exists(\"./logs\"):\n os.makedirs(\"./logs\")\n setlogger(args.log_file)\n\n # save the pre-trained model\n if not os.path.exists(\"./checkpoints\"):\n os.makedirs(\"./checkpoints\")\n\n # save the args\n for k, v in args.__dict__.items():\n logging.info(\"{}: {}\".format(k, v))\n \n trainer(args)" }, { "path": "loss/CORAL.py", "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\nimport torch\n\n\n\"\"\"\nCreated on Saturday Feb 25 2020\n\n@authors: Alan Preciado, Santosh Muthireddy\n\"\"\"\ndef CORAL_loss(source, target):\n\t\"\"\"\n\tFrom the paper, the vectors that compose Ds and Dt are D-dimensional vectors\n\t:param source: torch tensor: source data (Ds) with dimensions DxNs\n\t:param target: torch tensor: target data (Dt) with dimensons DxNt\n\t\"\"\"\n\n\td = source.size(1) # d-dimensional vectors (same for source, target)\n\n\tsource_covariance = compute_covariance(source)\n\ttarget_covariance = compute_covariance(target)\n\n\t# take Frobenius norm (https://pytorch.org/docs/stable/torch.html)\n\tloss = torch.norm(torch.mul((source_covariance-target_covariance),\n\t\t\t\t\t\t\t\t(source_covariance-target_covariance)), p=\"fro\")\n\n\t# loss = torch.norm(torch.mm((source_covariance-target_covariance),\n\t# \t\t\t\t\t\t\t(source_covariance-target_covariance)), p=\"fro\")\n\n\tloss = loss/(4*d*d)\n\n\treturn loss\n\n\ndef compute_covariance(data):\n\t\"\"\"\n\tCompute covariance matrix for given dataset as shown in paper (eqs 2 and 3).\n\t:param data: torch tensor: input source/target data\n\t\"\"\"\n\n\t# data dimensions: nxd (this for Ns or Nt)\n\tn = data.size(0) # get batch size\n\t#print(\"compute covariance bath size n:\", n)\n\n # check gpu or cpu support\n\tif data.is_cuda:\n\t\tdevice = torch.device(\"cuda\")\n\telse:\n\t\tdevice = torch.device(\"cpu\")\n\n\t# proper matrix multiplication for right side of equation (2)\n\tones_vector = torch.ones(n).resize(1, n).to(device=device) # 1xN dimensional vector (transposed)\n\tone_onto_D = torch.mm(ones_vector, data)\n\tmult_right_terms = torch.mm(one_onto_D.t(), one_onto_D)\n\tmult_right_terms = torch.div(mult_right_terms, n) # element-wise divison\n\n\t# matrix multiplication for left side of equation (2)\n\tmult_left_terms = torch.mm(data.t(), data)\n\n\tcovariance_matrix= 1/(n-1) * torch.add(mult_left_terms,-1*(mult_right_terms))\n\n\treturn covariance_matrix\n" }, { "path": "loss/MKMMD.py", "content": "# source code: https://github.com/ZhaoZhibin/UDTL (it seems the link is not the original source of the MK-MMD code)\n# Params:\n# source: source data\n# target: target data\n# kernel_mul:\n# kernel_num: the number of kernel\n# fix_sigma:\n\nimport torch\n\ndef guassian_kernel(source, target, kernel_mul=2.0, kernel_num=5, fix_sigma=None):\n n_samples = int(source.size()[0])+int(target.size()[0]) # the number of source+target\n total = torch.cat([source, target], dim=0)\n total0 = total.unsqueeze(0).expand(int(total.size(0)), int(total.size(0)), int(total.size(1)))\n total1 = total.unsqueeze(1).expand(int(total.size(0)), int(total.size(0)), int(total.size(1)))\n L2_distance = ((total0-total1)**2).sum(2)\n if fix_sigma:\n bandwidth = fix_sigma\n else:\n bandwidth = torch.sum(L2_distance.data) / (n_samples**2-n_samples)\n bandwidth /= kernel_mul ** (kernel_num // 2)\n bandwidth_list = [bandwidth * (kernel_mul**i) for i in range(kernel_num)]\n kernel_val = [torch.exp(-L2_distance / bandwidth_temp) for bandwidth_temp in bandwidth_list]\n return sum(kernel_val)#/len(kernel_val)\n\n\ndef MKMMD(source, target, kernel_mul=2.0, kernel_num=5, fix_sigma=None):\n batch_size = int(source.size()[0])\n kernels = guassian_kernel(source, target,\n kernel_mul=kernel_mul, kernel_num=kernel_num, fix_sigma=fix_sigma)\n XX = kernels[:batch_size, :batch_size]\n YY = kernels[batch_size:, batch_size:]\n XY = kernels[:batch_size, batch_size:]\n YX = kernels[batch_size:, :batch_size]\n loss = torch.mean(XX + YY - XY - YX)\n return loss\n" }, { "path": "loss/MMDLinear.py", "content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\nimport torch\n\n\n\"\"\"\nCreated on Sunday 22 Mar 2020\n\n@authors: Alan Preciado, Santosh Muthireddy\n\"\"\"\ndef MMDLinear(source_activation, target_activation):\n\t\"\"\"\n\tFrom the paper, the loss used is the maximum mean discrepancy (MMD)\n\t:param source: torch tensor: source data (Ds) with dimensions DxNs\n\t:param target: torch tensor: target data (Dt) with dimensons DxNt\n\t\"\"\"\n\n\tdiff_domains = source_activation - target_activation\n\tloss = torch.mean(torch.mm(diff_domains, torch.transpose(diff_domains, 0, 1)))\n\n\treturn loss\n" }, { "path": "loss/__init__.py", "content": "" } ]