完成模型日志记录模块

train/VGG16Net.py

@@ -1,5 +1,13 @@
+'''
+@作者:你遇了我321640253@qq.com
+@文件:VGG16Net.py
+@创建时间:2023 11 20
+
+模型网络结构
+'''
+import torch
 from torch import nn
-from torch.utils.data import DataLoader
+from torchsummary import summary
 
 # 定义 VGG16 网络结构
 class VGG16(nn.Module):
@@ -52,3 +60,10 @@ class VGG16(nn.Module):
       x = self.fc1(x)
       return x
 
+
+
+
+def getSummary(size:tuple):
+    model = VGG16()
+    model = model.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
+    summary(model, size)