def test_lstm(): class Net(nn.Module): def __init__(self, vocab_size=20, embed_dim=300, hidden_dim=512, num_layers=2): super().__init__() self.hidden_dim = hidden_dim self.embedding = nn.Embedding(vocab_size, embed_dim) self.encoder = nn.LSTM(embed_dim, hidden_dim, num_layers=num_layers) self.decoder = nn.Linear(hidden_dim, vocab_size) def forward(self, x): embed = self.embedding(x) out, hidden = self.encoder(embed) out = self.decoder(out) out = out.view(-1, out.size(2)) return out, hidden inp = torch.zeros((1, 100)).long() # [length, batch_size] df, df_total = summary(Net(), inp)
def test_multi_inputs(): class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1_x = nn.Conv2d(1, 10, kernel_size=5) self.conv2_x = nn.Conv2d(10, 20, kernel_size=5) self.conv2_drop_x = nn.Dropout2d() self.conv1_y = nn.Conv2d(1, 10, kernel_size=5) self.conv2_y = nn.Conv2d(10, 20, kernel_size=5) self.conv2_drop_y = nn.Dropout2d() self.fc1 = nn.Linear(320, 50) self.fc2 = nn.Linear(50, 10) def forward(self, x, y): y = F.relu(F.max_pool2d(self.conv1_y(y), 2)) y = F.relu(F.max_pool2d(self.conv2_drop_y(self.conv2_y(y)), 2)) x = F.relu(F.max_pool2d(self.conv1_x(x), 2)) x = F.relu(F.max_pool2d(self.conv2_drop_x(self.conv2_x(x)), 2)) x = x.view(-1, 320) x = F.relu(self.fc1(x)) x = F.dropout(x, training=self.training) x = self.fc2(x) return F.log_softmax(x, dim=1) inp1 = torch.zeros((1, 1, 28, 28)) inp2 = torch.zeros((1, 1, 64, 64)) summary_value, _ = summary(Net(), inp1, inp2)
def __init__(self, cfg): super(VAE, self).__init__(cfg) # Build Encoder self._encoder_construct() # Build Decoder self._decoder_construct() # Model Summary model_summary, _ = summary(self, torch.zeros((2, self.cfg.input_dim))) logging.info(f"====Model Summary====\n{model_summary}\n")
def __init__(self, cfg): super(VAE, self).__init__(cfg) self.latent1_dim = 20 self.in_dim = 32 self.layer1_dim = 64 self.layer2_dim = 128 self.latent2_dim = 20 # Build Encoder self._encoder_construct() # Build Decoder self._decoder_construct() # Model Summary model_summary, _ = summary(self, torch.zeros((2, self.cfg.input_dim))) logging.info(f"====Model Summary====\n{model_summary}\n")
def test_conv(): class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = nn.Conv2d(1, 10, kernel_size=5) self.conv2 = nn.Conv2d(10, 20, kernel_size=5) self.conv2_drop = nn.Dropout2d() self.fc1 = nn.Linear(320, 50) self.fc2 = nn.Linear(50, 10) def forward(self, x): x = F.relu(F.max_pool2d(self.conv1(x), 2)) x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) x = x.view(-1, 320) x = F.relu(self.fc1(x)) x = F.dropout(x, training=self.training) x = self.fc2(x) return F.log_softmax(x, dim=1) summary_value, _ = summary(Net(), torch.zeros((1, 1, 28, 28)))
def test_resnet_50(): model = torchvision.models.resnet50() summary_value, _ = summary(model, torch.zeros((4, 3, 224, 224)))
x = self.conv5(x) x = self.last(x) return x def resnet18(in_c, n_cls): return ResNet(in_c, n_cls, [2, 2, 2, 2], BasicBlock) def resnet34(in_c, n_cls): return ResNet(in_c, n_cls, [3, 4, 6, 3], BasicBlock) def resnet50(in_c, n_cls): return ResNet(in_c, n_cls, [3, 4, 6, 3]) def resnet101(in_c, n_cls): return ResNet(in_c, n_cls, [3, 4, 23, 3]) def resnet152(in_c, n_cls): return ResNet(in_c, n_cls, [3, 8, 36, 3]) if __name__ == "__main__": from torchvision.models import resnet from torchsummaryM import summary test_model = resnet18(3, 1000) official_model = resnet.resnet18() summary(test_model, torch.zeros(1, 3, 224, 224)) summary(official_model, torch.zeros(1, 3, 224, 224))
x = self.classifier(x) return x """ Define VGG16, VGG16_BN, VGG19, VGG19_BN """ def vgg16(in_c, n_cls): return VGG(in_c, n_cls, [2, 2, 3, 3, 3], _bn=False) def vgg19(in_c, n_cls): return VGG(in_c, n_cls, [2, 2, 4, 4, 4], _bn=False) def vgg16_bn(in_c, n_cls): return VGG(in_c, n_cls, [2, 2, 3, 3, 3], _bn=True) def vgg19_bn(in_c, n_cls): return VGG(in_c, n_cls, [2, 2, 4, 4, 4], _bn=True) if __name__ == "__main__": from torchsummaryM import summary from torchvision.models import vgg official_model = vgg.vgg16() test_model = vgg16(3, 1000) inps = torch.zeros(1, 3, 224, 224) summary(official_model, inps) summary(test_model, inps)
self.optim.load_state_dict(optim_dict) try: self.module.load_state_dict(state_dict) except AttributeError as e: self.load_state_dict(state_dict) self.cfg.start_epoch = epoch logging.info(f"{ckpt} model loaded.") logging.info(f"Epoch re-start at {epoch}.\n") else: # raise FileNotFoundError("File doesn't exsist. Please check the directory.") logging.info("File doesn't exsist. Please check the directory.") logging.info("Start training with initailized model.\n") if __name__ == "__main__": from pathlib import Path from torchsummaryM import summary class config: def __init__(self): # model self.latent_dim = 16 self.encode_hidden_dims = [64, 128] self.decode_hidden_dims = [64, 16] cfg = config() model = VAE(cfg) summary(model, torch.zeros((2, 16)))
size, baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, thickness=2) text_w, text_h = size x, y = left, top x1y1 = (x, y) x2y2 = (x + text_w + line_thickness, y + text_h + line_thickness + baseline) cv2.rectangle(image_boxes, x1y1, x2y2, bgr, -1) cv2.putText(image_boxes, text, (x + line_thickness, y + 2 * baseline + line_thickness), cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.4, color=(255, 255, 255), thickness=1, lineType=8) return image_boxes if __name__ == "__main__": from torchsummaryM import summary from torchvision.models import vgg16_bn inps = torch.ones(1, 3, 448, 448) yolo = YOLO_v1(_model='resnet') summary(yolo, inps)
def main(): args = parser.parse_args() device = "cuda" if args.cuda else "cpu" if device == "cuda": torch.backends.cudnn.benchmark = True # Set train & valid dataset dataset = VOC_Dataset(args.img, args.label, debug=args.debug) valid_ratio = .01 dataset_size = len(dataset) valid_size = int(valid_ratio * dataset_size) train_set, valid_set = random_split(dataset, [dataset_size-valid_size, valid_size]) train_loader = DataLoader(train_set, batch_size=args.batch, shuffle=True) valid_loader = DataLoader(valid_set, batch_size=1) # Define feature_extractor and Yolo model based on # [You Only Look Once: Unified, Real-Time Object Detection] https://arxiv.org/abs/1506.02640 yolo_model = YOLO_v1(bn=args.bn, _model=args.feature, _device=device).to(device) # Print summary of the model if needed if args.summary: summary(yolo_model, torch.ones(1, 3, 448, 448).to(device)) # Define criterion & optimizer based on # [You Only Look Once: Unified, Real-Time Object Detection] https://arxiv.org/abs/1506.02640 criterion = YOLO_loss(f_size=yolo_model.S, _gpu=True) optimizer = torch.optim.SGD(yolo_model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay) debug_dir = 'debug' os.makedirs(debug_dir, exist_ok=True) for epoch in range(args.start_epoch, args.epochs+1): yolo_model.train() cnt = 1 train_loss = 0 train_batch = 0 valid_loss = 0 valid_batch = 0 dataset._train = True for i, (inps, tars) in enumerate(train_loader): update_lr(optimizer, epoch, float(i) / float(len(train_loader) - 1), args=args) lr = get_lr(optimizer) inps = inps.float().to(device) tars = tars.float().to(device) preds = yolo_model(inps) loss = criterion(preds, tars) batch_size = inps.size(0) loss_cur_iter = loss.item() train_loss += loss_cur_iter * batch_size train_batch += batch_size optimizer.zero_grad() loss.backward() optimizer.step() print('Epoch [{:3d}/{:3d}] Iter [{:5d}/{:5d}] LR: {:.4f}'.format( epoch, args.epochs, i, len(train_loader), lr, end='')) print(' Loss: {:.4f} Average Loss: {:.4f}'.format( loss.item(), train_loss / float(train_batch) )) if i % 200 == 0: yolo_model.eval() dataset._train = False if i == 0: dataset._debug = True else: dataset._debug = False for j, (inps, tars) in enumerate(valid_loader): # Load data as a batch. inps = inps.float().to(device) # Forward to compute validation loss. with torch.no_grad(): boxes, class_names, probs = yolo_model.detect(image_bgr=inps, image_size=448, valid=True) img = (inps.squeeze(0).detach().cpu().numpy().transpose(1, 2, 0) * 127.5) + 127.5 plot_img = visualize_boxes(img, boxes, class_names, probs, valid=False) saveD = os.path.join(debug_dir, 'train_valid', '{:s}'.format(str(epoch).zfill(3))) os.makedirs(saveD, exist_ok=True) saveD = os.path.join(saveD, '{:s}'.format(str(i).zfill(5))) os.makedirs(saveD, exist_ok=True) cv2.imwrite(os.path.join(saveD, 'debug_{:s}.jpg'.format(str(cnt).zfill(5))), plot_img) cnt += 1 yolo_model.train() torch.save(yolo_model.state_dict(), "./yolo_v1.ptr")
Inverted_Residual_Block(160, 320, s=1, t=6, alpha=alpha), nn.Conv2d(int(320 * alpha), int(alpha * 1280), 1, 1, 0, bias=False), nn.BatchNorm2d(int(alpha * 1280)), nn.ReLU6(True) ] feature = nn.Sequential(*layers) return feature def _build_classifier(self, alpha): layers = [ nn.AdaptiveAvgPool2d(output_size=(1, 1)), nn.Flatten(start_dim=1), nn.Linear(in_features=int(alpha * 1280), out_features=1000), ] classifier = nn.Sequential(*layers) return classifier if __name__ == "__main__": from torchsummaryM import summary from torchvision.models.mobilenet import MobileNetV2 model = MobileNetV2() summary(model, torch.zeros(1, 3, 224, 224)) model = MobileNet_v2(alpha=1.4) summary(model, torch.zeros(1, 3, 224, 224))