查看网络结构

最新推荐文章于 2023-03-14 18:24:08 发布

yyfhq

最新推荐文章于 2023-03-14 18:24:08 发布

阅读量433

点赞数

分类专栏： pytorch 文章标签： python

本文链接：https://blog.csdn.net/yyfhq/article/details/124127008

版权

pytorch 专栏收录该内容

8 篇文章 0 订阅

订阅专栏

import torch
from torchsummary import summary

from modelsun import SunNet

if __name__ == "__main__":
    # 需要使用device来指定网络在GPU还是CPU运行
    device  = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model   = SunNet(num_classes=6, aux_logits=False, init_weights=False).to(device)
    summary(model, (3,224,224))