root@node2:/home/usr/resnet50# python resnet50-rand.py
get hyhal driver version error, open file(/sys/module/hydcu/version and /sys/module/amdgpu/version) error, skip init hydmi
Using device: cuda:0, cudnn enabled = False
Using model: Resnet50
Batch [1/10], Time: 4.489s, Loss: 4.8839,
Batch [2/10], Time: 1.209s, Loss: 4.6027,
Batch [3/10], Time: 1.208s, Loss: 4.3348,
Batch [4/10], Time: 1.198s, Loss: 4.2885,
Batch [5/10], Time: 1.206s, Loss: 4.3044,
Batch [6/10], Time: 1.196s, Loss: 4.3034,
Batch [7/10], Time: 1.208s, Loss: 4.3115,
Batch [8/10], Time: 1.209s, Loss: 4.3045,
Batch [9/10], Time: 1.216s, Loss: 4.2469,
Batch [10/10], Time: 1.205s, Loss: 4.2284,
total time: 16.754 s
如果能拉取镜像的话,试试下面这个镜像:docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.4.1-ubuntu22.04-dtk25.04-py3.10-fixpy(亲测可成功运行无上述报错)
以下是python源码:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import argparse
import time
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.models as models
import os
def parse_args():
parser = argparse.ArgumentParser(description="Train ResNet50 on randomly generated data")
parser.add_argument("--use_cudnn", type=bool, default=False, help="Whether to use cudnn (True/False)")
parser.add_argument("--batch_size", type=int, default=128, help="Batch size for training")
parser.add_argument("--num_batches", type=int, default=10, help="Number of batches to run")
parser.add_argument("--num_classes", type=int, default=100, help="Number of classes for the random labels")
return parser.parse_args()
def main():
os.environ["TORCH_CUDNN_V8_API_DISABLED"] = "1"
os.environ["PYTORCH_NO_CUDA_MEMORY_CACHING"] = "1"
args = parse_args()
# 配置 cudnn
torch.backends.cudnn.enabled = args.use_cudnn
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}, cudnn enabled = {args.use_cudnn}")
print(f"Using model: Resnet50")
model = models.resnet50(weights=None)
model.fc = nn.Linear(model.fc.in_features, 100)
model.to(device)
criterion = nn.CrossEntropyLoss().to(device)
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
inputs = torch.randn(args.batch_size, 3, 224, 224)
targets = torch.randint(0, args.num_classes, (args.batch_size,))
inputs = inputs.to(device)
targets = targets.to(device)
total_time_start=time.time()
for batch_idx in range(args.num_batches):
start_time = time.time()
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
end_time = time.time()
batch_time = end_time - start_time
print(
f"Batch [{batch_idx+1}/{args.num_batches}], "
f"Time: {batch_time:.3f}s, "
f"Loss: {loss.item():.4f}, "
)
total_time_end=time.time()
total_time=total_time_end - total_time_start
print(f"\ntotal time: {total_time:.3f} s")
if __name__ == "__main__":
main()