PyTorch 中 PyCharm 计算网络单集运行时间
class DQN(nn.Module): def init(self, in_planes, outputs, stride=1, hidden_size=128): super(DQN, self).init()
self.conv1 = nn.Conv2d(in_channels=in_planes, out_channels=256, kernel_size=3,
stride=stride, padding=1, bias=False)
self.conv2 = nn.Conv2d(in_channels=256, out_channels=128, kernel_size=3,
stride=stride, padding=1, bias=False)
self.conv3 = nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3,
stride=stride, padding=1, bias=False)
# self.fc1 = nn.Linear(1568,64)
################################################ ## 残差+ se注意力 ## ################################################ self.conv1 = self._make_layer(SEBasicBlock, in_planes, 256, 2, stride=stride) self.conv2 = self._make_layer(SEBasicBlock, 256, 128, 2, stride=stride) self.conv3 = self._make_layer(SEBasicBlock, 128, 64, 2, stride=stride)
# self.fc2 = nn.Linear(64, 128)
# self.fc3 = nn.Linear(128, 256)
# self.val_hidden = nn.Linear(256, 128)
# self.adv_hidden = nn.Linear(256, 128)
# self.val = nn.Linear(128, 1)
# self.adv = nn.Linear(128, outputs)
start_time=time.time()
self.fc1 = nn.Linear(64*4*4, 128)
self.fc2 = nn.Linear(128, 256)
self.val_hidden = nn.Linear(256, 128)
self.adv_hidden = nn.Linear(256,128)
self.val = nn.Linear(128, 1)
self.adv = nn.Linear(128, outputs)
end_time=time.time()
print('网络结构加载时间:',end_time-start_time) # 添加输出语句
如何调用该网络结构
可以通过以下方式调用该网络结构:
- 创建一个DQN实例:
model = DQN(in_planes=3, outputs=10)
其中,in_planes表示输入图像的通道数,outputs表示网络输出的类别数。这里假设输入图像为RGB三通道,输出类别数为10。
- 将输入数据送入网络中进行计算:
input_data = torch.randn(1, 3, 32, 32) # 假设输入图像大小为32x32 output = model(input_data)
其中,input_data表示输入的数据,output表示网络的输出。这里假设输入数据为大小为1x3x32x32的张量。可以通过调用torch.randn函数生成一个随机的输入数据。
- 计算网络单集运行时间:
start_time = time.time() output = model(input_data) end_time = time.time() running_time = end_time - start_time print('网络单集运行时间: ', running_time)
代码示例:
import torch
import torch.nn as nn
import time
# 定义网络结构
class DQN(nn.Module):
# ... (代码同上)
# 创建网络实例
model = DQN(in_planes=3, outputs=10)
# 生成随机输入数据
input_data = torch.randn(1, 3, 32, 32)
# 计算网络单集运行时间
start_time = time.time()
output = model(input_data)
end_time = time.time()
running_time = end_time - start_time
print('网络单集运行时间: ', running_time)
运行上述代码即可得到网络单集运行时间。
原文地址: https://www.cveoy.top/t/topic/m7qY 著作权归作者所有。请勿转载和采集!